Skip to content

Commit f9950ce

Browse files
dlarocquedlarocque
committed
feat(vertexai): Add support for AbortSignal 
1 parent 61f0102 commit f9950ce

15 files changed

+260
-60
lines changed

‎.changeset/khaki-ladybugs-kiss.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@firebase/vertexai': minor
3+
---
4+
5+
Add an `AbortSignal` to request options for generation methods, allowing them to be aborted.

‎common/api-review/vertexai.api.md

+12-7
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@ export class ChatSession {
5050
params?: StartChatParams | undefined;
5151
// (undocumented)
5252
requestOptions?: RequestOptions | undefined;
53-
sendMessage(request: string | Array<string | Part>): Promise<GenerateContentResult>;
54-
sendMessageStream(request: string | Array<string | Part>): Promise<GenerateContentStreamResult>;
53+
sendMessage(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
54+
sendMessageStream(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
5555
}
5656

5757
// @public
@@ -325,9 +325,9 @@ export interface GenerativeContentBlob {
325325
// @public
326326
export class GenerativeModel extends VertexAIModel {
327327
constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions);
328-
countTokens(request: CountTokensRequest | string | Array<string | Part>): Promise<CountTokensResponse>;
329-
generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;
330-
generateContentStream(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentStreamResult>;
328+
countTokens(request: CountTokensRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<CountTokensResponse>;
329+
generateContent(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
330+
generateContentStream(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
331331
// (undocumented)
332332
generationConfig: GenerationConfig;
333333
// (undocumented)
@@ -463,9 +463,9 @@ export interface ImagenInlineImage {
463463
// @beta
464464
export class ImagenModel extends VertexAIModel {
465465
constructor(vertexAI: VertexAI, modelParams: ImagenModelParams, requestOptions?: RequestOptions | undefined);
466-
generateImages(prompt: string): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
466+
generateImages(prompt: string, singleRequestOptions?: SingleRequestOptions): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
467467
// @internal
468-
generateImagesGCS(prompt: string, gcsURI: string): Promise<ImagenGenerationResponse<ImagenGCSImage>>;
468+
generateImagesGCS(prompt: string, gcsURI: string, singleRequestOptions?: SingleRequestOptions): Promise<ImagenGenerationResponse<ImagenGCSImage>>;
469469
generationConfig?: ImagenGenerationConfig;
470470
// (undocumented)
471471
requestOptions?: RequestOptions | undefined;
@@ -719,6 +719,11 @@ export interface Segment {
719719
startIndex: number;
720720
}
721721

722+
// @public
723+
export interface SingleRequestOptions extends RequestOptions {
724+
signal?: AbortSignal;
725+
}
726+
722727
// @public
723728
export interface StartChatParams extends BaseParams {
724729
// (undocumented)

‎docs-devsite/_toc.yaml

+2
Original file line numberDiff line numberDiff line change
@@ -588,6 +588,8 @@ toc:
588588
path: /docs/reference/js/vertexai.schemashared.md
589589
- title: Segment
590590
path: /docs/reference/js/vertexai.segment.md
591+
- title: SingleRequestOptions
592+
path: /docs/reference/js/vertexai.singlerequestoptions.md
591593
- title: StartChatParams
592594
path: /docs/reference/js/vertexai.startchatparams.md
593595
- title: StringSchema

‎docs-devsite/vertexai.chatsession.md

+6-4
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,8 @@ export declare class ChatSession
3737
| Method | Modifiers | Description |
3838
| --- | --- | --- |
3939
| [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. |
40-
| [sendMessage(request)](./vertexai.chatsession.md#chatsessionsendmessage) | | Sends a chat message and receives a non-streaming <code>[GenerateContentResult](./vertexai.generatecontentresult.md#generatecontentresult_interface)</code> |
41-
| [sendMessageStream(request)](./vertexai.chatsession.md#chatsessionsendmessagestream) | | Sends a chat message and receives the response as a <code>[GenerateContentStreamResult](./vertexai.generatecontentstreamresult.md#generatecontentstreamresult_interface)</code> containing an iterable stream and a response promise. |
40+
| [sendMessage(request, singleRequestOptions)](./vertexai.chatsession.md#chatsessionsendmessage) | | Sends a chat message and receives a non-streaming <code>[GenerateContentResult](./vertexai.generatecontentresult.md#generatecontentresult_interface)</code> |
41+
| [sendMessageStream(request, singleRequestOptions)](./vertexai.chatsession.md#chatsessionsendmessagestream) | | Sends a chat message and receives the response as a <code>[GenerateContentStreamResult](./vertexai.generatecontentstreamresult.md#generatecontentstreamresult_interface)</code> containing an iterable stream and a response promise. |
4242

4343
## ChatSession.(constructor)
4444

@@ -103,14 +103,15 @@ Sends a chat message and receives a non-streaming <code>[GenerateContentResult](
103103
<b>Signature:</b>
104104

105105
```typescript
106-
sendMessage(request: string | Array<string | Part>): Promise<GenerateContentResult>;
106+
sendMessage(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
107107
```
108108

109109
#### Parameters
110110

111111
| Parameter | Type | Description |
112112
| --- | --- | --- |
113113
| request | string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; | |
114+
| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | |
114115

115116
<b>Returns:</b>
116117

@@ -123,14 +124,15 @@ Sends a chat message and receives the response as a <code>[GenerateContentStream
123124
<b>Signature:</b>
124125

125126
```typescript
126-
sendMessageStream(request: string | Array<string | Part>): Promise<GenerateContentStreamResult>;
127+
sendMessageStream(request: string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
127128
```
128129

129130
#### Parameters
130131

131132
| Parameter | Type | Description |
132133
| --- | --- | --- |
133134
| request | string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; | |
135+
| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | |
134136

135137
<b>Returns:</b>
136138

‎docs-devsite/vertexai.generativemodel.md

+9-6
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,9 @@ export declare class GenerativeModel extends VertexAIModel
4040
4141
| Method | Modifiers | Description |
4242
| --- | --- | --- |
43-
| [countTokens(request)](./vertexai.generativemodel.md#generativemodelcounttokens) | | Counts the tokens in the provided request. |
44-
| [generateContent(request)](./vertexai.generativemodel.md#generativemodelgeneratecontent) | | Makes a single non-streaming call to the model and returns an object containing a single <code>[GenerateContentResponse](./vertexai.generatecontentresponse.md#generatecontentresponse_interface)</code>. |
45-
| [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. |
43+
| [countTokens(request, singleRequestOptions)](./vertexai.generativemodel.md#generativemodelcounttokens) | | Counts the tokens in the provided request. |
44+
| [generateContent(request, singleRequestOptions)](./vertexai.generativemodel.md#generativemodelgeneratecontent) | | Makes a single non-streaming call to the model and returns an object containing a single <code>[GenerateContentResponse](./vertexai.generatecontentresponse.md#generatecontentresponse_interface)</code>. |
45+
| [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. |
4646
| [startChat(startChatParams)](./vertexai.generativemodel.md#generativemodelstartchat) | | Gets a new <code>[ChatSession](./vertexai.chatsession.md#chatsession_class)</code> instance which can be used for multi-turn chats. |
4747
4848
## GenerativeModel.(constructor)
@@ -118,14 +118,15 @@ Counts the tokens in the provided request.
118118
<b>Signature:</b>
119119
120120
```typescript
121-
countTokens(request: CountTokensRequest | string | Array<string | Part>): Promise<CountTokensResponse>;
121+
countTokens(request: CountTokensRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<CountTokensResponse>;
122122
```
123123
124124
#### Parameters
125125
126126
| Parameter | Type | Description |
127127
| --- | --- | --- |
128128
| request | [CountTokensRequest](./vertexai.counttokensrequest.md#counttokensrequest_interface) \| string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; | |
129+
| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | |
129130
130131
<b>Returns:</b>
131132
@@ -138,14 +139,15 @@ Makes a single non-streaming call to the model and returns an object containing
138139
<b>Signature:</b>
139140
140141
```typescript
141-
generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;
142+
generateContent(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
142143
```
143144
144145
#### Parameters
145146
146147
| Parameter | Type | Description |
147148
| --- | --- | --- |
148149
| request | [GenerateContentRequest](./vertexai.generatecontentrequest.md#generatecontentrequest_interface) \| string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; | |
150+
| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | |
149151
150152
<b>Returns:</b>
151153
@@ -158,14 +160,15 @@ Makes a single streaming call to the model and returns an object containing an i
158160
<b>Signature:</b>
159161
160162
```typescript
161-
generateContentStream(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentStreamResult>;
163+
generateContentStream(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
162164
```
163165
164166
#### Parameters
165167
166168
| Parameter | Type | Description |
167169
| --- | --- | --- |
168170
| request | [GenerateContentRequest](./vertexai.generatecontentrequest.md#generatecontentrequest_interface) \| string \| Array&lt;string \| [Part](./vertexai.md#part)<!-- -->&gt; | |
171+
| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | |
169172
170173
<b>Returns:</b>
171174

‎docs-devsite/vertexai.imagenmodel.md

+3-2
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ export declare class ImagenModel extends VertexAIModel
4242
4343
| Method | Modifiers | Description |
4444
| --- | --- | --- |
45-
| [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. |
45+
| [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. |
4646
4747
## ImagenModel.(constructor)
4848
@@ -118,14 +118,15 @@ If the prompt was not blocked, but one or more of the generated images were filt
118118
<b>Signature:</b>
119119
120120
```typescript
121-
generateImages(prompt: string): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
121+
generateImages(prompt: string, singleRequestOptions?: SingleRequestOptions): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
122122
```
123123
124124
#### Parameters
125125
126126
| Parameter | Type | Description |
127127
| --- | --- | --- |
128128
| prompt | string | A text prompt describing the image(s) to generate. |
129+
| singleRequestOptions | [SingleRequestOptions](./vertexai.singlerequestoptions.md#singlerequestoptions_interface) | |
129130
130131
<b>Returns:</b>
131132

‎docs-devsite/vertexai.md

+1
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ The Vertex AI in Firebase Web SDK.
111111
| [SchemaRequest](./vertexai.schemarequest.md#schemarequest_interface) | Final format for <code>[Schema](./vertexai.schema.md#schema_class)</code> params passed to backend requests. |
112112
| [SchemaShared](./vertexai.schemashared.md#schemashared_interface) | Basic <code>[Schema](./vertexai.schema.md#schema_class)</code> properties shared across several Schema-related types. |
113113
| [Segment](./vertexai.segment.md#segment_interface) | |
114+
| [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)<!-- -->). |
114115
| [StartChatParams](./vertexai.startchatparams.md#startchatparams_interface) | Params for [GenerativeModel.startChat()](./vertexai.generativemodel.md#generativemodelstartchat)<!-- -->. |
115116
| [TextPart](./vertexai.textpart.md#textpart_interface) | Content part interface if the part represents a text string. |
116117
| [ToolConfig](./vertexai.toolconfig.md#toolconfig_interface) | Tool config. This config is shared for all tools provided in the request. |

‎docs-devsite/vertexai.singlerequestoptions.md

+57
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
Project: /docs/reference/js/_project.yaml
2+
Book: /docs/reference/_book.yaml
3+
page_type: reference
4+
5+
{% comment %}
6+
DO NOT EDIT THIS FILE!
7+
This is generated by the JS SDK team, and any local changes will be
8+
overwritten. Changes should be made in the source code at
9+
https://github.com/firebase/firebase-js-sdk
10+
{% endcomment %}
11+
12+
# SingleRequestOptions interface
13+
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`<!-- -->.
14+
15+
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)<!-- -->).
16+
17+
<b>Signature:</b>
18+
19+
```typescript
20+
export interface SingleRequestOptions extends RequestOptions
21+
```
22+
<b>Extends:</b> [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface)
23+
24+
## Properties
25+
26+
| Property | Type | Description |
27+
| --- | --- | --- |
28+
| [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. |
29+
30+
## SingleRequestOptions.signal
31+
32+
An `AbortSignal` instance that allows cancelling ongoing requests (like `generateContent` or `generateImages`<!-- -->).
33+
34+
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.
35+
36+
Note that this will not cancel the request in the backend, so billing will still be applied despite cancellation.
37+
38+
<b>Signature:</b>
39+
40+
```typescript
41+
signal?: AbortSignal;
42+
```
43+
44+
### Example
45+
46+
47+
```javascript
48+
const controller = new AbortController();
49+
const model = getGenerativeModel({
50+
// ...
51+
});
52+
53+
// To cancel request:
54+
controller.abort();
55+
56+
```
57+

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

+15-4
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import {
2222
GenerateContentStreamResult,
2323
Part,
2424
RequestOptions,
25+
SingleRequestOptions,
2526
StartChatParams
2627
} from '../types';
2728
import { formatNewContent } from '../requests/request-helpers';
@@ -75,7 +76,8 @@ export class ChatSession {
7576
* <code>{@link GenerateContentResult}</code>
7677
*/
7778
async sendMessage(
78-
request: string | Array<string | Part>
79+
request: string | Array<string | Part>,
80+
singleRequestOptions?: SingleRequestOptions
7981
): Promise<GenerateContentResult> {
8082
await this._sendPromise;
8183
const newContent = formatNewContent(request);
@@ -95,7 +97,11 @@ export class ChatSession {
9597
this._apiSettings,
9698
this.model,
9799
generateContentRequest,
98-
this.requestOptions
100+
// Merge requestOptions
101+
{
102+
...singleRequestOptions,
103+
...this.requestOptions
104+
}
99105
)
100106
)
101107
.then(result => {
@@ -130,7 +136,8 @@ export class ChatSession {
130136
* and a response promise.
131137
*/
132138
async sendMessageStream(
133-
request: string | Array<string | Part>
139+
request: string | Array<string | Part>,
140+
singleRequestOptions?: SingleRequestOptions
134141
): Promise<GenerateContentStreamResult> {
135142
await this._sendPromise;
136143
const newContent = formatNewContent(request);
@@ -146,7 +153,11 @@ export class ChatSession {
146153
this._apiSettings,
147154
this.model,
148155
generateContentRequest,
149-
this.requestOptions
156+
// Merge requestOptions
157+
{
158+
...singleRequestOptions,
159+
...this.requestOptions
160+
}
150161
);
151162

152163
// Add onto the chain.

‎packages/vertexai/src/methods/count-tokens.ts

+3-3
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
import {
1919
CountTokensRequest,
2020
CountTokensResponse,
21-
RequestOptions
21+
SingleRequestOptions
2222
} from '../types';
2323
import { Task, makeRequest } from '../requests/request';
2424
import { ApiSettings } from '../types/internal';
@@ -27,15 +27,15 @@ export async function countTokens(
2727
apiSettings: ApiSettings,
2828
model: string,
2929
params: CountTokensRequest,
30-
requestOptions?: RequestOptions
30+
singleRequestOptions?: SingleRequestOptions
3131
): Promise<CountTokensResponse> {
3232
const response = await makeRequest(
3333
model,
3434
Task.COUNT_TOKENS,
3535
apiSettings,
3636
false,
3737
JSON.stringify(params),
38-
requestOptions
38+
singleRequestOptions
3939
);
4040
return response.json();
4141
}

‎packages/vertexai/src/methods/generate-content.ts

+5-5
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import {
2020
GenerateContentResponse,
2121
GenerateContentResult,
2222
GenerateContentStreamResult,
23-
RequestOptions
23+
SingleRequestOptions
2424
} from '../types';
2525
import { Task, makeRequest } from '../requests/request';
2626
import { createEnhancedContentResponse } from '../requests/response-helpers';
@@ -31,15 +31,15 @@ export async function generateContentStream(
3131
apiSettings: ApiSettings,
3232
model: string,
3333
params: GenerateContentRequest,
34-
requestOptions?: RequestOptions
34+
singleRequestOptions?: SingleRequestOptions
3535
): Promise<GenerateContentStreamResult> {
3636
const response = await makeRequest(
3737
model,
3838
Task.STREAM_GENERATE_CONTENT,
3939
apiSettings,
4040
/* stream */ true,
4141
JSON.stringify(params),
42-
requestOptions
42+
singleRequestOptions
4343
);
4444
return processStream(response);
4545
}
@@ -48,15 +48,15 @@ export async function generateContent(
4848
apiSettings: ApiSettings,
4949
model: string,
5050
params: GenerateContentRequest,
51-
requestOptions?: RequestOptions
51+
singleRequestOptions?: SingleRequestOptions
5252
): Promise<GenerateContentResult> {
5353
const response = await makeRequest(
5454
model,
5555
Task.GENERATE_CONTENT,
5656
apiSettings,
5757
/* stream */ false,
5858
JSON.stringify(params),
59-
requestOptions
59+
singleRequestOptions
6060
);
6161
const responseJson: GenerateContentResponse = await response.json();
6262
const enhancedResponse = createEnhancedContentResponse(responseJson);

0 commit comments

Comments
 (0)