Skip to content

Commit bf1b031

Browse files
dlarocquedlarocque
committed
throw AbortError if timeout is invoked on request
1 parent 40ce3c5 commit bf1b031

File tree

10 files changed

+211
-38
lines changed

10 files changed

+211
-38
lines changed

‎.vscode/launch.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
"src/index.node.ts",
3838
"--timeout",
3939
"5000",
40-
"integration/**/*.test.ts"
40+
"integration/**/prompt-templates.test.ts"
4141
],
4242
"env": {
4343
"TS_NODE_COMPILER_OPTIONS": "{\"module\":\"commonjs\"}"

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1338,8 +1338,9 @@ export class TemplateGenerativeModel {
13381338
constructor(ai: AI, requestOptions?: RequestOptions);
13391339
// @internal (undocumented)
13401340
_apiSettings: ApiSettings;
1341-
generateContent(templateId: string, templateVariables: object): Promise<GenerateContentResult>;
1342-
generateContentStream(templateId: string, templateVariables: object): Promise<GenerateContentStreamResult>;
1341+
generateContent(templateId: string, templateVariables: object, // anything!
1342+
singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentResult>;
1343+
generateContentStream(templateId: string, templateVariables: object, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
13431344
requestOptions?: RequestOptions;
13441345
}
13451346

@@ -1348,7 +1349,7 @@ export class TemplateImagenModel {
13481349
constructor(ai: AI, requestOptions?: RequestOptions);
13491350
// @internal (undocumented)
13501351
_apiSettings: ApiSettings;
1351-
generateImages(templateId: string, templateVariables: object): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
1352+
generateImages(templateId: string, templateVariables: object, singleRequestOptions?: SingleRequestOptions): Promise<ImagenGenerationResponse<ImagenInlineImage>>;
13521353
requestOptions?: RequestOptions;
13531354
}
13541355

‎packages/ai/integration/prompt-templates.test.ts‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -35,16 +35,25 @@ describe('Prompt templates', function () {
3535
describe(`${testConfig.toString()}`, () => {
3636
describe('Generative Model', () => {
3737
it('successfully generates content', async () => {
38+
const a = new AbortController();
3839
const model = getTemplateGenerativeModel(testConfig.ai, {
3940
baseUrl: STAGING_URL
4041
});
41-
const { response } = await model.generateContent(
42-
`sassy-greeting-${templateBackendSuffix(
43-
testConfig.ai.backend.backendType
44-
)}`,
45-
{ name: 'John' }
46-
);
47-
expect(response.text()).to.contain('John'); // Template asks to address directly by name
42+
// a.abort();
43+
try {
44+
await model.generateContent(
45+
`sassy-greeting-${templateBackendSuffix(
46+
testConfig.ai.backend.backendType
47+
)}`,
48+
{ name: 'John' },
49+
{ signal: a.signal, timeout: 100 }
50+
);
51+
} catch (e) {
52+
console.error(e);
53+
if ((e as DOMException).name === 'AbortError') {
54+
console.log(1);
55+
}
56+
}
4857
});
4958
});
5059
describe('Imagen model', async () => {
@@ -56,7 +65,8 @@ describe('Prompt templates', function () {
5665
`portrait-${templateBackendSuffix(
5766
testConfig.ai.backend.backendType
5867
)}`,
59-
{ animal: 'Rhino' }
68+
{ animal: 'Rhino' },
69+
{ timeout: 100}
6070
);
6171
expect(images.length).to.equal(2); // We ask for two images in the prompt template
6272
});

‎packages/ai/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
"test:browser": "yarn testsetup && karma start",
4242
"test:node": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' mocha --require ts-node/register --require src/index.node.ts 'src/**/!(*-browser)*.test.ts' --config ../../config/mocharc.node.js",
4343
"test:integration": "karma start --integration",
44-
"test:integration:node": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' mocha integration/**/*.test.ts --config ../../config/mocharc.node.js",
44+
"test:integration:node": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' mocha integration/**/prompt-templates.test.ts --config ../../config/mocharc.node.js",
4545
"api-report": "api-extractor run --local --verbose",
4646
"typings:public": "node ../../scripts/build/use_typings.js ./dist/ai-public.d.ts",
4747
"type-check": "yarn tsc --noEmit",

‎packages/ai/src/models/template-generative-model.test.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,51 @@ describe('TemplateGenerativeModel', () => {
7373
{ timeout: 5000 }
7474
);
7575
});
76+
77+
it('singleRequestOptions overrides requestOptions', async () => {
78+
const templateGenerateContentStub = stub(
79+
generateContentMethods,
80+
'templateGenerateContent'
81+
).resolves({} as any);
82+
const model = new TemplateGenerativeModel(fakeAI, { timeout: 1000 });
83+
const singleRequestOptions = { timeout: 2000 };
84+
85+
await model.generateContent(
86+
TEMPLATE_ID,
87+
TEMPLATE_VARS,
88+
singleRequestOptions
89+
);
90+
91+
expect(templateGenerateContentStub).to.have.been.calledOnceWith(
92+
model._apiSettings,
93+
TEMPLATE_ID,
94+
{ inputs: TEMPLATE_VARS },
95+
{ timeout: 2000 }
96+
);
97+
});
98+
99+
it('singleRequestOptions is merged with requestOptions', async () => {
100+
const templateGenerateContentStub = stub(
101+
generateContentMethods,
102+
'templateGenerateContent'
103+
).resolves({} as any);
104+
const abortController = new AbortController();
105+
const model = new TemplateGenerativeModel(fakeAI, { timeout: 1000 });
106+
const singleRequestOptions = { signal: abortController.signal };
107+
108+
await model.generateContent(
109+
TEMPLATE_ID,
110+
TEMPLATE_VARS,
111+
singleRequestOptions
112+
);
113+
114+
expect(templateGenerateContentStub).to.have.been.calledOnceWith(
115+
model._apiSettings,
116+
TEMPLATE_ID,
117+
{ inputs: TEMPLATE_VARS },
118+
{ timeout: 1000, signal: abortController.signal }
119+
);
120+
});
76121
});
77122

78123
describe('generateContentStream', () => {
@@ -92,5 +137,50 @@ describe('TemplateGenerativeModel', () => {
92137
{ timeout: 5000 }
93138
);
94139
});
140+
141+
it('singleRequestOptions overrides requestOptions', async () => {
142+
const templateGenerateContentStreamStub = stub(
143+
generateContentMethods,
144+
'templateGenerateContentStream'
145+
).resolves({} as any);
146+
const model = new TemplateGenerativeModel(fakeAI, { timeout: 1000 });
147+
const singleRequestOptions = { timeout: 2000 };
148+
149+
await model.generateContentStream(
150+
TEMPLATE_ID,
151+
TEMPLATE_VARS,
152+
singleRequestOptions
153+
);
154+
155+
expect(templateGenerateContentStreamStub).to.have.been.calledOnceWith(
156+
model._apiSettings,
157+
TEMPLATE_ID,
158+
{ inputs: TEMPLATE_VARS },
159+
{ timeout: 2000 }
160+
);
161+
});
162+
163+
it('singleRequestOptions is merged with requestOptions', async () => {
164+
const templateGenerateContentStreamStub = stub(
165+
generateContentMethods,
166+
'templateGenerateContentStream'
167+
).resolves({} as any);
168+
const abortController = new AbortController();
169+
const model = new TemplateGenerativeModel(fakeAI, { timeout: 1000 });
170+
const singleRequestOptions = { signal: abortController.signal };
171+
172+
await model.generateContentStream(
173+
TEMPLATE_ID,
174+
TEMPLATE_VARS,
175+
singleRequestOptions
176+
);
177+
178+
expect(templateGenerateContentStreamStub).to.have.been.calledOnceWith(
179+
model._apiSettings,
180+
TEMPLATE_ID,
181+
{ inputs: TEMPLATE_VARS },
182+
{ timeout: 1000, signal: abortController.signal }
183+
);
184+
});
95185
});
96186
});

‎packages/ai/src/models/template-generative-model.ts‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,11 @@ import {
2020
templateGenerateContentStream
2121
} from '../methods/generate-content';
2222
import { GenerateContentResult, RequestOptions } from '../types';
23-
import { AI, GenerateContentStreamResult } from '../public-types';
23+
import {
24+
AI,
25+
GenerateContentStreamResult,
26+
SingleRequestOptions
27+
} from '../public-types';
2428
import { ApiSettings } from '../types/internal';
2529
import { initApiSettings } from './utils';
2630

@@ -62,13 +66,17 @@ export class TemplateGenerativeModel {
6266
*/
6367
async generateContent(
6468
templateId: string,
65-
templateVariables: object // anything!
69+
templateVariables: object, // anything!
70+
singleRequestOptions?: SingleRequestOptions
6671
): Promise<GenerateContentResult> {
6772
return templateGenerateContent(
6873
this._apiSettings,
6974
templateId,
7075
{ inputs: templateVariables },
71-
this.requestOptions // TODO: Add singleRequestOptions parameter and merge both request options here.
76+
{
77+
...this.requestOptions,
78+
...singleRequestOptions
79+
}
7280
);
7381
}
7482

@@ -86,13 +94,17 @@ export class TemplateGenerativeModel {
8694
*/
8795
async generateContentStream(
8896
templateId: string,
89-
templateVariables: object
97+
templateVariables: object,
98+
singleRequestOptions?: SingleRequestOptions
9099
): Promise<GenerateContentStreamResult> {
91100
return templateGenerateContentStream(
92101
this._apiSettings,
93102
templateId,
94103
{ inputs: templateVariables },
95-
this.requestOptions // TODO: Add singleRequestOptions parameter and merge both request options here.
104+
{
105+
...this.requestOptions,
106+
...singleRequestOptions
107+
}
96108
);
97109
}
98110
}

‎packages/ai/src/models/template-imagen-model.test.ts‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
import { use, expect } from 'chai';
1919
import sinonChai from 'sinon-chai';
2020
import chaiAsPromised from 'chai-as-promised';
21-
import { restore, stub } from 'sinon';
21+
import { restore, stub, match } from 'sinon';
2222
import { AI } from '../public-types';
2323
import { VertexAIBackend } from '../backend';
2424
import { TemplateImagenModel } from './template-imagen-model';
@@ -89,6 +89,62 @@ describe('TemplateImagenModel', () => {
8989
);
9090
});
9191

92+
it('singleRequestOptions overrides requestOptions', async () => {
93+
const mockPrediction = {
94+
'bytesBase64Encoded':
95+
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==',
96+
'mimeType': 'image/png'
97+
};
98+
const makeRequestStub = stub(request, 'makeRequest').resolves({
99+
json: () => Promise.resolve({ predictions: [mockPrediction] })
100+
} as Response);
101+
const model = new TemplateImagenModel(fakeAI, { timeout: 1000 });
102+
const singleRequestOptions = { timeout: 2000 };
103+
104+
await model.generateImages(
105+
TEMPLATE_ID,
106+
TEMPLATE_VARS,
107+
singleRequestOptions
108+
);
109+
110+
expect(makeRequestStub).to.have.been.calledOnceWith(
111+
match({
112+
singleRequestOptions: { timeout: 2000 }
113+
}),
114+
match.any
115+
);
116+
});
117+
118+
it('singleRequestOptions is merged with requestOptions', async () => {
119+
const mockPrediction = {
120+
'bytesBase64Encoded':
121+
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==',
122+
'mimeType': 'image/png'
123+
};
124+
const makeRequestStub = stub(request, 'makeRequest').resolves({
125+
json: () => Promise.resolve({ predictions: [mockPrediction] })
126+
} as Response);
127+
const abortController = new AbortController();
128+
const model = new TemplateImagenModel(fakeAI, { timeout: 1000 });
129+
const singleRequestOptions = { signal: abortController.signal };
130+
131+
await model.generateImages(
132+
TEMPLATE_ID,
133+
TEMPLATE_VARS,
134+
singleRequestOptions
135+
);
136+
137+
expect(makeRequestStub).to.have.been.calledOnceWith(
138+
match({
139+
singleRequestOptions: {
140+
timeout: 1000,
141+
signal: abortController.signal
142+
}
143+
}),
144+
match.any
145+
);
146+
});
147+
92148
it('should return the result of handlePredictResponse', async () => {
93149
const mockPrediction = {
94150
'bytesBase64Encoded':

‎packages/ai/src/models/template-imagen-model.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ import { RequestOptions } from '../types';
1919
import {
2020
AI,
2121
ImagenGenerationResponse,
22-
ImagenInlineImage
22+
ImagenInlineImage,
23+
SingleRequestOptions
2324
} from '../public-types';
2425
import { ApiSettings } from '../types/internal';
2526
import { makeRequest, ServerPromptTemplateTask } from '../requests/request';
@@ -64,15 +65,19 @@ export class TemplateImagenModel {
6465
*/
6566
async generateImages(
6667
templateId: string,
67-
templateVariables: object
68+
templateVariables: object,
69+
singleRequestOptions?: SingleRequestOptions
6870
): Promise<ImagenGenerationResponse<ImagenInlineImage>> {
6971
const response = await makeRequest(
7072
{
7173
task: ServerPromptTemplateTask.TEMPLATE_PREDICT,
7274
templateId,
7375
apiSettings: this._apiSettings,
7476
stream: false,
75-
singleRequestOptions: this.requestOptions // TODO: Add singleRequestOptions parameter and merge both request options here.
77+
singleRequestOptions: {
78+
...this.requestOptions,
79+
...singleRequestOptions
80+
}
7681
},
7782
JSON.stringify({ inputs: templateVariables })
7883
);

0 commit comments

Comments
 (0)