Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit b7a3d24

Browse files
gsiddherikeldridge
gsiddh
and
erikeldridge
committedApr 22, 2025
Define ChromeAdapter class (#8942)
Co-authored-by: Erik Eldridge <[email protected]>
1 parent eeeab4e commit b7a3d24

File tree

11 files changed

+256
-78
lines changed

11 files changed

+256
-78
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
static DEFAULT_HYBRID_IN_CLOUD_MODEL: string;
330331
generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;

‎docs-devsite/vertexai.chatsession.md

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ export declare class ChatSession
2222

2323
| Constructor | Modifiers | Description |
2424
| --- | --- | --- |
25-
| [(constructor)(apiSettings, model, params, requestOptions)](./vertexai.chatsession.md#chatsessionconstructor) | | Constructs a new instance of the <code>ChatSession</code> class |
25+
| [(constructor)(apiSettings, model, chromeAdapter, params, requestOptions)](./vertexai.chatsession.md#chatsessionconstructor) | | Constructs a new instance of the <code>ChatSession</code> class |
2626

2727
## Properties
2828

@@ -47,7 +47,7 @@ Constructs a new instance of the `ChatSession` class
4747
<b>Signature:</b>
4848

4949
```typescript
50-
constructor(apiSettings: ApiSettings, model: string, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
50+
constructor(apiSettings: ApiSettings, model: string, chromeAdapter: ChromeAdapter, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
5151
```
5252

5353
#### Parameters
@@ -56,6 +56,7 @@ constructor(apiSettings: ApiSettings, model: string, params?: StartChatParams |
5656
| --- | --- | --- |
5757
| apiSettings | ApiSettings | |
5858
| model | string | |
59+
| chromeAdapter | ChromeAdapter | |
5960
| params | [StartChatParams](./vertexai.startchatparams.md#startchatparams_interface) \| undefined | |
6061
| requestOptions | [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) \| undefined | |
6162

‎docs-devsite/vertexai.generativemodel.md

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ export declare class GenerativeModel extends VertexAIModel
2323
2424
| Constructor | Modifiers | Description |
2525
| --- | --- | --- |
26-
| [(constructor)(vertexAI, modelParams, requestOptions)](./vertexai.generativemodel.md#generativemodelconstructor) | | Constructs a new instance of the <code>GenerativeModel</code> class |
26+
| [(constructor)(vertexAI, modelParams, chromeAdapter, requestOptions)](./vertexai.generativemodel.md#generativemodelconstructor) | | Constructs a new instance of the <code>GenerativeModel</code> class |
2727
2828
## Properties
2929
@@ -53,7 +53,7 @@ Constructs a new instance of the `GenerativeModel` class
5353
<b>Signature:</b>
5454
5555
```typescript
56-
constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions);
56+
constructor(vertexAI: VertexAI, modelParams: ModelParams, chromeAdapter: ChromeAdapter, requestOptions?: RequestOptions);
5757
```
5858
5959
#### Parameters
@@ -62,6 +62,7 @@ constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: Reque
6262
| --- | --- | --- |
6363
| vertexAI | [VertexAI](./vertexai.vertexai.md#vertexai_interface) | |
6464
| modelParams | [ModelParams](./vertexai.modelparams.md#modelparams_interface) | |
65+
| chromeAdapter | ChromeAdapter | |
6566
| requestOptions | [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) | |
6667
6768
## GenerativeModel.DEFAULT\_HYBRID\_IN\_CLOUD\_MODEL

‎packages/vertexai/src/api.ts

Lines changed: 7 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,12 @@ export function getGenerativeModel(
9192
`Must provide a model name. Example: getGenerativeModel({ model: 'my-model-name' })`
9293
);
9394
}
94-
return new GenerativeModel(vertexAI, inCloudParams, requestOptions);
95+
return new GenerativeModel(
96+
vertexAI,
97+
inCloudParams,
98+
new ChromeAdapter(hybridParams.mode, hybridParams.onDeviceParams),
99+
requestOptions
100+
);
95101
}
96102

97103
/**

‎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: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
import { GenerateContentRequest, InferenceMode } from '../types';
19+
import { LanguageModelCreateOptions } from '../types/language-model';
20+
21+
/**
22+
* Defines an inference "backend" that uses Chrome's on-device model,
23+
* and encapsulates logic for detecting when on-device is possible.
24+
*/
25+
export class ChromeAdapter {
26+
constructor(
27+
private mode?: InferenceMode,
28+
private onDeviceParams?: LanguageModelCreateOptions
29+
) {}
30+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
31+
async isAvailable(request: GenerateContentRequest): Promise<boolean> {
32+
return false;
33+
}
34+
async generateContentOnDevice(
35+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
36+
request: GenerateContentRequest
37+
): Promise<Response> {
38+
return {
39+
json: () =>
40+
Promise.resolve({
41+
candidates: [
42+
{
43+
content: {
44+
parts: [{ text: '' }]
45+
}
46+
}
47+
]
48+
})
49+
} as Response;
50+
}
51+
}

‎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);
@@ -70,7 +71,8 @@ describe('generateContent()', () => {
7071
const result = await generateContent(
7172
fakeApiSettings,
7273
'model',
73-
fakeRequestParams
74+
fakeRequestParams,
75+
new ChromeAdapter()
7476
);
7577
expect(result.response.text()).to.include('Mountain View, California');
7678
expect(makeRequestStub).to.be.calledWith(
@@ -95,7 +97,8 @@ describe('generateContent()', () => {
9597
const result = await generateContent(
9698
fakeApiSettings,
9799
'model',
98-
fakeRequestParams
100+
fakeRequestParams,
101+
new ChromeAdapter()
99102
);
100103
expect(result.response.text()).to.include('Use Freshly Ground Coffee');
101104
expect(result.response.text()).to.include('30 minutes of brewing');
@@ -118,7 +121,8 @@ describe('generateContent()', () => {
118121
const result = await generateContent(
119122
fakeApiSettings,
120123
'model',
121-
fakeRequestParams
124+
fakeRequestParams,
125+
new ChromeAdapter()
122126
);
123127
expect(result.response.usageMetadata?.totalTokenCount).to.equal(1913);
124128
expect(result.response.usageMetadata?.candidatesTokenCount).to.equal(76);
@@ -153,7 +157,8 @@ describe('generateContent()', () => {
153157
const result = await generateContent(
154158
fakeApiSettings,
155159
'model',
156-
fakeRequestParams
160+
fakeRequestParams,
161+
new ChromeAdapter()
157162
);
158163
expect(result.response.text()).to.include(
159164
'Some information cited from an external source'
@@ -180,7 +185,8 @@ describe('generateContent()', () => {
180185
const result = await generateContent(
181186
fakeApiSettings,
182187
'model',
183-
fakeRequestParams
188+
fakeRequestParams,
189+
new ChromeAdapter()
184190
);
185191
expect(result.response.text).to.throw('SAFETY');
186192
expect(makeRequestStub).to.be.calledWith(
@@ -202,7 +208,8 @@ describe('generateContent()', () => {
202208
const result = await generateContent(
203209
fakeApiSettings,
204210
'model',
205-
fakeRequestParams
211+
fakeRequestParams,
212+
new ChromeAdapter()
206213
);
207214
expect(result.response.text).to.throw('SAFETY');
208215
expect(makeRequestStub).to.be.calledWith(
@@ -224,7 +231,8 @@ describe('generateContent()', () => {
224231
const result = await generateContent(
225232
fakeApiSettings,
226233
'model',
227-
fakeRequestParams
234+
fakeRequestParams,
235+
new ChromeAdapter()
228236
);
229237
expect(result.response.text()).to.equal('');
230238
expect(makeRequestStub).to.be.calledWith(
@@ -246,7 +254,8 @@ describe('generateContent()', () => {
246254
const result = await generateContent(
247255
fakeApiSettings,
248256
'model',
249-
fakeRequestParams
257+
fakeRequestParams,
258+
new ChromeAdapter()
250259
);
251260
expect(result.response.text()).to.include('Some text');
252261
expect(makeRequestStub).to.be.calledWith(
@@ -268,7 +277,12 @@ describe('generateContent()', () => {
268277
json: mockResponse.json
269278
} as Response);
270279
await expect(
271-
generateContent(fakeApiSettings, 'model', fakeRequestParams)
280+
generateContent(
281+
fakeApiSettings,
282+
'model',
283+
fakeRequestParams,
284+
new ChromeAdapter()
285+
)
272286
).to.be.rejectedWith(/400.*invalid argument/);
273287
expect(mockFetch).to.be.called;
274288
});
@@ -283,10 +297,36 @@ describe('generateContent()', () => {
283297
json: mockResponse.json
284298
} as Response);
285299
await expect(
286-
generateContent(fakeApiSettings, 'model', fakeRequestParams)
300+
generateContent(
301+
fakeApiSettings,
302+
'model',
303+
fakeRequestParams,
304+
new ChromeAdapter()
305+
)
287306
).to.be.rejectedWith(
288307
/firebasevertexai\.googleapis[\s\S]*my-project[\s\S]*api-not-enabled/
289308
);
290309
expect(mockFetch).to.be.called;
291310
});
311+
it('on-device', async () => {
312+
const chromeAdapter = new ChromeAdapter();
313+
const isAvailableStub = stub(chromeAdapter, 'isAvailable').resolves(true);
314+
const mockResponse = getMockResponse(
315+
'vertexAI',
316+
'unary-success-basic-reply-short.json'
317+
);
318+
const generateContentStub = stub(
319+
chromeAdapter,
320+
'generateContentOnDevice'
321+
).resolves(mockResponse as Response);
322+
const result = await generateContent(
323+
fakeApiSettings,
324+
'model',
325+
fakeRequestParams,
326+
chromeAdapter
327+
);
328+
expect(result.response.text()).to.include('Mountain View, California');
329+
expect(isAvailableStub).to.be.called;
330+
expect(generateContentStub).to.be.calledWith(fakeRequestParams);
331+
});
292332
});

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

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import { Task, makeRequest } from '../requests/request';
2626
import { createEnhancedContentResponse } from '../requests/response-helpers';
2727
import { processStream } from '../requests/stream-reader';
2828
import { ApiSettings } from '../types/internal';
29+
import { ChromeAdapter } from './chrome-adapter';
2930

3031
export async function generateContentStream(
3132
apiSettings: ApiSettings,
@@ -44,20 +45,40 @@ export async function generateContentStream(
4445
return processStream(response);
4546
}
4647

47-
export async function generateContent(
48+
async function generateContentOnCloud(
4849
apiSettings: ApiSettings,
4950
model: string,
5051
params: GenerateContentRequest,
5152
requestOptions?: RequestOptions
52-
): Promise<GenerateContentResult> {
53-
const response = await makeRequest(
53+
): Promise<Response> {
54+
return makeRequest(
5455
model,
5556
Task.GENERATE_CONTENT,
5657
apiSettings,
5758
/* stream */ false,
5859
JSON.stringify(params),
5960
requestOptions
6061
);
62+
}
63+
64+
export async function generateContent(
65+
apiSettings: ApiSettings,
66+
model: string,
67+
params: GenerateContentRequest,
68+
chromeAdapter: ChromeAdapter,
69+
requestOptions?: RequestOptions
70+
): Promise<GenerateContentResult> {
71+
let response;
72+
if (await chromeAdapter.isAvailable(params)) {
73+
response = await chromeAdapter.generateContentOnDevice(params);
74+
} else {
75+
response = await generateContentOnCloud(
76+
apiSettings,
77+
model,
78+
params,
79+
requestOptions
80+
);
81+
}
6182
const responseJson: GenerateContentResponse = await response.json();
6283
const enhancedResponse = createEnhancedContentResponse(responseJson);
6384
return {

‎packages/vertexai/src/models/generative-model.test.ts

Lines changed: 92 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import * as request from '../requests/request';
2121
import { match, restore, stub } from 'sinon';
2222
import { getMockResponse } from '../../test-utils/mock-response';
2323
import sinonChai from 'sinon-chai';
24+
import { ChromeAdapter } from '../methods/chrome-adapter';
2425

2526
use(sinonChai);
2627

@@ -39,21 +40,27 @@ const fakeVertexAI: VertexAI = {
3940

4041
describe('GenerativeModel', () => {
4142
it('passes params through to generateContent', async () => {
42-
const genModel = new GenerativeModel(fakeVertexAI, {
43-
model: 'my-model',
44-
tools: [
45-
{
46-
functionDeclarations: [
47-
{
48-
name: 'myfunc',
49-
description: 'mydesc'
50-
}
51-
]
52-
}
53-
],
54-
toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
55-
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
56-
});
43+
const genModel = new GenerativeModel(
44+
fakeVertexAI,
45+
{
46+
model: 'my-model',
47+
tools: [
48+
{
49+
functionDeclarations: [
50+
{
51+
name: 'myfunc',
52+
description: 'mydesc'
53+
}
54+
]
55+
}
56+
],
57+
toolConfig: {
58+
functionCallingConfig: { mode: FunctionCallingMode.NONE }
59+
},
60+
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
61+
},
62+
new ChromeAdapter()
63+
);
5764
expect(genModel.tools?.length).to.equal(1);
5865
expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
5966
FunctionCallingMode.NONE
@@ -84,10 +91,14 @@ describe('GenerativeModel', () => {
8491
restore();
8592
});
8693
it('passes text-only systemInstruction through to generateContent', async () => {
87-
const genModel = new GenerativeModel(fakeVertexAI, {
88-
model: 'my-model',
89-
systemInstruction: 'be friendly'
90-
});
94+
const genModel = new GenerativeModel(
95+
fakeVertexAI,
96+
{
97+
model: 'my-model',
98+
systemInstruction: 'be friendly'
99+
},
100+
new ChromeAdapter()
101+
);
91102
expect(genModel.systemInstruction?.parts[0].text).to.equal('be friendly');
92103
const mockResponse = getMockResponse(
93104
'vertexAI',
@@ -110,21 +121,27 @@ describe('GenerativeModel', () => {
110121
restore();
111122
});
112123
it('generateContent overrides model values', async () => {
113-
const genModel = new GenerativeModel(fakeVertexAI, {
114-
model: 'my-model',
115-
tools: [
116-
{
117-
functionDeclarations: [
118-
{
119-
name: 'myfunc',
120-
description: 'mydesc'
121-
}
122-
]
123-
}
124-
],
125-
toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
126-
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
127-
});
124+
const genModel = new GenerativeModel(
125+
fakeVertexAI,
126+
{
127+
model: 'my-model',
128+
tools: [
129+
{
130+
functionDeclarations: [
131+
{
132+
name: 'myfunc',
133+
description: 'mydesc'
134+
}
135+
]
136+
}
137+
],
138+
toolConfig: {
139+
functionCallingConfig: { mode: FunctionCallingMode.NONE }
140+
},
141+
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
142+
},
143+
new ChromeAdapter()
144+
);
128145
expect(genModel.tools?.length).to.equal(1);
129146
expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
130147
FunctionCallingMode.NONE
@@ -166,14 +183,20 @@ describe('GenerativeModel', () => {
166183
restore();
167184
});
168185
it('passes params through to chat.sendMessage', async () => {
169-
const genModel = new GenerativeModel(fakeVertexAI, {
170-
model: 'my-model',
171-
tools: [
172-
{ functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
173-
],
174-
toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
175-
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
176-
});
186+
const genModel = new GenerativeModel(
187+
fakeVertexAI,
188+
{
189+
model: 'my-model',
190+
tools: [
191+
{ functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
192+
],
193+
toolConfig: {
194+
functionCallingConfig: { mode: FunctionCallingMode.NONE }
195+
},
196+
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
197+
},
198+
new ChromeAdapter()
199+
);
177200
expect(genModel.tools?.length).to.equal(1);
178201
expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
179202
FunctionCallingMode.NONE
@@ -204,10 +227,14 @@ describe('GenerativeModel', () => {
204227
restore();
205228
});
206229
it('passes text-only systemInstruction through to chat.sendMessage', async () => {
207-
const genModel = new GenerativeModel(fakeVertexAI, {
208-
model: 'my-model',
209-
systemInstruction: 'be friendly'
210-
});
230+
const genModel = new GenerativeModel(
231+
fakeVertexAI,
232+
{
233+
model: 'my-model',
234+
systemInstruction: 'be friendly'
235+
},
236+
new ChromeAdapter()
237+
);
211238
expect(genModel.systemInstruction?.parts[0].text).to.equal('be friendly');
212239
const mockResponse = getMockResponse(
213240
'vertexAI',
@@ -230,14 +257,20 @@ describe('GenerativeModel', () => {
230257
restore();
231258
});
232259
it('startChat overrides model values', async () => {
233-
const genModel = new GenerativeModel(fakeVertexAI, {
234-
model: 'my-model',
235-
tools: [
236-
{ functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
237-
],
238-
toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
239-
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
240-
});
260+
const genModel = new GenerativeModel(
261+
fakeVertexAI,
262+
{
263+
model: 'my-model',
264+
tools: [
265+
{ functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
266+
],
267+
toolConfig: {
268+
functionCallingConfig: { mode: FunctionCallingMode.NONE }
269+
},
270+
systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
271+
},
272+
new ChromeAdapter()
273+
);
241274
expect(genModel.tools?.length).to.equal(1);
242275
expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
243276
FunctionCallingMode.NONE
@@ -282,7 +315,11 @@ describe('GenerativeModel', () => {
282315
restore();
283316
});
284317
it('calls countTokens', async () => {
285-
const genModel = new GenerativeModel(fakeVertexAI, { model: 'my-model' });
318+
const genModel = new GenerativeModel(
319+
fakeVertexAI,
320+
{ model: 'my-model' },
321+
new ChromeAdapter()
322+
);
286323
const mockResponse = getMockResponse(
287324
'vertexAI',
288325
'unary-success-total-tokens.json'

‎packages/vertexai/src/models/generative-model.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ import {
4343
} from '../requests/request-helpers';
4444
import { VertexAI } from '../public-types';
4545
import { VertexAIModel } from './vertexai-model';
46+
import { ChromeAdapter } from '../methods/chrome-adapter';
4647

4748
/**
4849
* Class for generative model APIs.
@@ -63,6 +64,7 @@ export class GenerativeModel extends VertexAIModel {
6364
constructor(
6465
vertexAI: VertexAI,
6566
modelParams: ModelParams,
67+
private chromeAdapter: ChromeAdapter,
6668
requestOptions?: RequestOptions
6769
) {
6870
super(vertexAI, modelParams.model);
@@ -95,6 +97,7 @@ export class GenerativeModel extends VertexAIModel {
9597
systemInstruction: this.systemInstruction,
9698
...formattedParams
9799
},
100+
this.chromeAdapter,
98101
this.requestOptions
99102
);
100103
}
@@ -132,6 +135,7 @@ export class GenerativeModel extends VertexAIModel {
132135
return new ChatSession(
133136
this._apiSettings,
134137
this.model,
138+
this.chromeAdapter,
135139
{
136140
tools: this.tools,
137141
toolConfig: this.toolConfig,

0 commit comments

Comments
 (0)
Please sign in to comment.