Skip to content

Commit c891a85

Browse files
erikeldridgeerikeldridge
committed
Define ChromeAdapter class
1 parent 2c8915e commit c891a85

File tree

10 files changed

+396
-73
lines changed

10 files changed

+396
-73
lines changed

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ export class BooleanSchema extends Schema {
4242
// @public
4343
export class ChatSession {
4444
// Warning: (ae-forgotten-export) The symbol "ApiSettings" needs to be exported by the entry point index.d.ts
45-
constructor(apiSettings: ApiSettings, model: string, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
45+
// Warning: (ae-forgotten-export) The symbol "ChromeAdapter" needs to be exported by the entry point index.d.ts
46+
constructor(apiSettings: ApiSettings, model: string, chromeAdapter: ChromeAdapter, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
4647
getHistory(): Promise<Content[]>;
4748
// (undocumented)
4849
model: string;
@@ -324,7 +325,7 @@ export interface GenerativeContentBlob {
324325

325326
// @public
326327
export class GenerativeModel extends VertexAIModel {
327-
constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions);
328+
constructor(vertexAI: VertexAI, modelParams: ModelParams, chromeAdapter: ChromeAdapter, requestOptions?: RequestOptions);
328329
countTokens(request: CountTokensRequest | string | Array<string | Part>): Promise<CountTokensResponse>;
329330
generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;
330331
generateContentStream(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentStreamResult>;

‎packages/vertexai/src/api.ts‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import {
3030
} from './types';
3131
import { VertexAIError } from './errors';
3232
import { VertexAIModel, GenerativeModel, ImagenModel } from './models';
33+
import { ChromeAdapter } from './methods/chrome-adapter';
3334

3435
export { ChatSession } from './methods/chat-session';
3536
export * from './requests/schema-builder';
@@ -91,7 +92,16 @@ export function getGenerativeModel(
9192
`Must provide a model name. Example: getGenerativeModel({ model: 'my-model-name' })`
9293
);
9394
}
94-
return new GenerativeModel(vertexAI, onCloudParams, requestOptions);
95+
return new GenerativeModel(
96+
vertexAI,
97+
onCloudParams,
98+
new ChromeAdapter(
99+
window.ai as AI,
100+
hybridParams.mode,
101+
hybridParams.onDeviceParams
102+
),
103+
requestOptions
104+
);
95105
}
96106

97107
/**

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import * as generateContentMethods from './generate-content';
2323
import { GenerateContentStreamResult } from '../types';
2424
import { ChatSession } from './chat-session';
2525
import { ApiSettings } from '../types/internal';
26+
import { ChromeAdapter } from './chrome-adapter';
2627

2728
use(sinonChai);
2829
use(chaiAsPromised);
@@ -44,7 +45,11 @@ describe('ChatSession', () => {
4445
generateContentMethods,
4546
'generateContent'
4647
).rejects('generateContent failed');
47-
const chatSession = new ChatSession(fakeApiSettings, 'a-model');
48+
const chatSession = new ChatSession(
49+
fakeApiSettings,
50+
'a-model',
51+
new ChromeAdapter()
52+
);
4853
await expect(chatSession.sendMessage('hello')).to.be.rejected;
4954
expect(generateContentStub).to.be.calledWith(
5055
fakeApiSettings,
@@ -61,7 +66,11 @@ describe('ChatSession', () => {
6166
generateContentMethods,
6267
'generateContentStream'
6368
).rejects('generateContentStream failed');
64-
const chatSession = new ChatSession(fakeApiSettings, 'a-model');
69+
const chatSession = new ChatSession(
70+
fakeApiSettings,
71+
'a-model',
72+
new ChromeAdapter()
73+
);
6574
await expect(chatSession.sendMessageStream('hello')).to.be.rejected;
6675
expect(generateContentStreamStub).to.be.calledWith(
6776
fakeApiSettings,
@@ -80,7 +89,11 @@ describe('ChatSession', () => {
8089
generateContentMethods,
8190
'generateContentStream'
8291
).resolves({} as unknown as GenerateContentStreamResult);
83-
const chatSession = new ChatSession(fakeApiSettings, 'a-model');
92+
const chatSession = new ChatSession(
93+
fakeApiSettings,
94+
'a-model',
95+
new ChromeAdapter()
96+
);
8497
await chatSession.sendMessageStream('hello');
8598
expect(generateContentStreamStub).to.be.calledWith(
8699
fakeApiSettings,

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import { validateChatHistory } from './chat-session-helpers';
3030
import { generateContent, generateContentStream } from './generate-content';
3131
import { ApiSettings } from '../types/internal';
3232
import { logger } from '../logger';
33+
import { ChromeAdapter } from './chrome-adapter';
3334

3435
/**
3536
* Do not log a message for this error.
@@ -50,6 +51,7 @@ export class ChatSession {
5051
constructor(
5152
apiSettings: ApiSettings,
5253
public model: string,
54+
private chromeAdapter: ChromeAdapter,
5355
public params?: StartChatParams,
5456
public requestOptions?: RequestOptions
5557
) {
@@ -95,6 +97,7 @@ export class ChatSession {
9597
this._apiSettings,
9698
this.model,
9799
generateContentRequest,
100+
this.chromeAdapter,
98101
this.requestOptions
99102
)
100103
)

‎packages/vertexai/src/methods/chrome-adapter.ts‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
import {
2+
EnhancedGenerateContentResponse,
3+
GenerateContentRequest,
4+
InferenceMode
5+
} from '../types';
6+
import { AI, AILanguageModelCreateOptionsWithSystemPrompt } from '../types/ai';
7+
8+
/**
9+
* Defines an inference "backend" that uses Chrome's on-device model,
10+
* and encapsulates logic for detecting when on-device is possible.
11+
*/
12+
export class ChromeAdapter {
13+
constructor(
14+
private aiProvider?: AI,
15+
private mode?: InferenceMode,
16+
private onDeviceParams?: AILanguageModelCreateOptionsWithSystemPrompt
17+
) {}
18+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
19+
async isAvailable(request: GenerateContentRequest): Promise<boolean> {
20+
return false;
21+
}
22+
async generateContentOnDevice(
23+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
24+
request: GenerateContentRequest
25+
): Promise<EnhancedGenerateContentResponse> {
26+
return {
27+
text: () => '',
28+
functionCalls: () => undefined
29+
};
30+
}
31+
}

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

Lines changed: 50 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import {
3030
} from '../types';
3131
import { ApiSettings } from '../types/internal';
3232
import { Task } from '../requests/request';
33+
import { ChromeAdapter } from './chrome-adapter';
3334

3435
use(sinonChai);
3536
use(chaiAsPromised);
@@ -69,7 +70,8 @@ describe('generateContent()', () => {
6970
const result = await generateContent(
7071
fakeApiSettings,
7172
'model',
72-
fakeRequestParams
73+
fakeRequestParams,
74+
new ChromeAdapter()
7375
);
7476
expect(result.response.text()).to.include('Mountain View, California');
7577
expect(makeRequestStub).to.be.calledWith(
@@ -91,7 +93,8 @@ describe('generateContent()', () => {
9193
const result = await generateContent(
9294
fakeApiSettings,
9395
'model',
94-
fakeRequestParams
96+
fakeRequestParams,
97+
new ChromeAdapter()
9598
);
9699
expect(result.response.text()).to.include('Use Freshly Ground Coffee');
97100
expect(result.response.text()).to.include('30 minutes of brewing');
@@ -113,7 +116,8 @@ describe('generateContent()', () => {
113116
const result = await generateContent(
114117
fakeApiSettings,
115118
'model',
116-
fakeRequestParams
119+
fakeRequestParams,
120+
new ChromeAdapter()
117121
);
118122
expect(result.response.usageMetadata?.totalTokenCount).to.equal(1913);
119123
expect(result.response.usageMetadata?.candidatesTokenCount).to.equal(76);
@@ -145,7 +149,8 @@ describe('generateContent()', () => {
145149
const result = await generateContent(
146150
fakeApiSettings,
147151
'model',
148-
fakeRequestParams
152+
fakeRequestParams,
153+
new ChromeAdapter()
149154
);
150155
expect(result.response.text()).to.include(
151156
'Some information cited from an external source'
@@ -171,7 +176,8 @@ describe('generateContent()', () => {
171176
const result = await generateContent(
172177
fakeApiSettings,
173178
'model',
174-
fakeRequestParams
179+
fakeRequestParams,
180+
new ChromeAdapter()
175181
);
176182
expect(result.response.text).to.throw('SAFETY');
177183
expect(makeRequestStub).to.be.calledWith(
@@ -192,7 +198,8 @@ describe('generateContent()', () => {
192198
const result = await generateContent(
193199
fakeApiSettings,
194200
'model',
195-
fakeRequestParams
201+
fakeRequestParams,
202+
new ChromeAdapter()
196203
);
197204
expect(result.response.text).to.throw('SAFETY');
198205
expect(makeRequestStub).to.be.calledWith(
@@ -211,7 +218,8 @@ describe('generateContent()', () => {
211218
const result = await generateContent(
212219
fakeApiSettings,
213220
'model',
214-
fakeRequestParams
221+
fakeRequestParams,
222+
new ChromeAdapter()
215223
);
216224
expect(result.response.text()).to.equal('');
217225
expect(makeRequestStub).to.be.calledWith(
@@ -232,7 +240,8 @@ describe('generateContent()', () => {
232240
const result = await generateContent(
233241
fakeApiSettings,
234242
'model',
235-
fakeRequestParams
243+
fakeRequestParams,
244+
new ChromeAdapter()
236245
);
237246
expect(result.response.text()).to.include('Some text');
238247
expect(makeRequestStub).to.be.calledWith(
@@ -251,7 +260,12 @@ describe('generateContent()', () => {
251260
json: mockResponse.json
252261
} as Response);
253262
await expect(
254-
generateContent(fakeApiSettings, 'model', fakeRequestParams)
263+
generateContent(
264+
fakeApiSettings,
265+
'model',
266+
fakeRequestParams,
267+
new ChromeAdapter()
268+
)
255269
).to.be.rejectedWith(/400.*invalid argument/);
256270
expect(mockFetch).to.be.called;
257271
});
@@ -265,10 +279,36 @@ describe('generateContent()', () => {
265279
json: mockResponse.json
266280
} as Response);
267281
await expect(
268-
generateContent(fakeApiSettings, 'model', fakeRequestParams)
282+
generateContent(
283+
fakeApiSettings,
284+
'model',
285+
fakeRequestParams,
286+
new ChromeAdapter()
287+
)
269288
).to.be.rejectedWith(
270289
/firebasevertexai\.googleapis[\s\S]*my-project[\s\S]*api-not-enabled/
271290
);
272291
expect(mockFetch).to.be.called;
273292
});
293+
it('on-device', async () => {
294+
const expectedText = 'hi';
295+
const chromeAdapter = new ChromeAdapter();
296+
const mockIsAvailable = stub(chromeAdapter, 'isAvailable').resolves(true);
297+
const mockGenerateContent = stub(
298+
chromeAdapter,
299+
'generateContentOnDevice'
300+
).resolves({
301+
text: () => expectedText,
302+
functionCalls: () => undefined
303+
});
304+
const result = await generateContent(
305+
fakeApiSettings,
306+
'model',
307+
fakeRequestParams,
308+
chromeAdapter
309+
);
310+
expect(result.response.text()).to.equal(expectedText);
311+
expect(mockIsAvailable).to.be.called;
312+
expect(mockGenerateContent).to.be.calledWith(fakeRequestParams);
313+
});
274314
});

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

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
*/
1717

1818
import {
19+
EnhancedGenerateContentResponse,
1920
GenerateContentRequest,
2021
GenerateContentResponse,
2122
GenerateContentResult,
@@ -26,6 +27,7 @@ import { Task, makeRequest } from '../requests/request';
2627
import { createEnhancedContentResponse } from '../requests/response-helpers';
2728
import { processStream } from '../requests/stream-reader';
2829
import { ApiSettings } from '../types/internal';
30+
import { ChromeAdapter } from './chrome-adapter';
2931

3032
export async function generateContentStream(
3133
apiSettings: ApiSettings,
@@ -44,12 +46,12 @@ export async function generateContentStream(
4446
return processStream(response);
4547
}
4648

47-
export async function generateContent(
49+
async function generateContentOnCloud(
4850
apiSettings: ApiSettings,
4951
model: string,
5052
params: GenerateContentRequest,
5153
requestOptions?: RequestOptions
52-
): Promise<GenerateContentResult> {
54+
): Promise<EnhancedGenerateContentResponse> {
5355
const response = await makeRequest(
5456
model,
5557
Task.GENERATE_CONTENT,
@@ -60,6 +62,27 @@ export async function generateContent(
6062
);
6163
const responseJson: GenerateContentResponse = await response.json();
6264
const enhancedResponse = createEnhancedContentResponse(responseJson);
65+
return enhancedResponse;
66+
}
67+
68+
export async function generateContent(
69+
apiSettings: ApiSettings,
70+
model: string,
71+
params: GenerateContentRequest,
72+
chromeAdapter: ChromeAdapter,
73+
requestOptions?: RequestOptions
74+
): Promise<GenerateContentResult> {
75+
let enhancedResponse;
76+
if (await chromeAdapter.isAvailable(params)) {
77+
enhancedResponse = await chromeAdapter.generateContentOnDevice(params);
78+
} else {
79+
enhancedResponse = await generateContentOnCloud(
80+
apiSettings,
81+
model,
82+
params,
83+
requestOptions
84+
);
85+
}
6386
return {
6487
response: enhancedResponse
6588
};

0 commit comments

Comments
 (0)