Skip to content

Commit fb13a48

Browse files
erikeldridgegsiddh
erikeldridge
authored and
gsiddh
committed
VinF Hybrid Inference: Implement ChromeAdapter (rebased) (#8943)
1 parent 28bf975 commit fb13a48

File tree

7 files changed

+514
-19
lines changed

7 files changed

+514
-19
lines changed

‎e2e/sample-apps/modular.js

+6-1
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,12 @@ import {
5858
onValue,
5959
off
6060
} from 'firebase/database';
61-
import { getGenerativeModel, getVertexAI, VertexAI } from 'firebase/vertexai';
61+
import {
62+
getGenerativeModel,
63+
getVertexAI,
64+
InferenceMode,
65+
VertexAI
66+
} from 'firebase/vertexai';
6267
import { getDataConnect, DataConnect } from 'firebase/data-connect';
6368

6469
/**

‎packages/vertexai/src/api.ts

+6-1
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import {
3131
import { VertexAIError } from './errors';
3232
import { VertexAIModel, GenerativeModel, ImagenModel } from './models';
3333
import { ChromeAdapter } from './methods/chrome-adapter';
34+
import { LanguageModel } from './types/language-model';
3435

3536
export { ChatSession } from './methods/chat-session';
3637
export * from './requests/schema-builder';
@@ -95,7 +96,11 @@ export function getGenerativeModel(
9596
return new GenerativeModel(
9697
vertexAI,
9798
inCloudParams,
98-
new ChromeAdapter(hybridParams.mode, hybridParams.onDeviceParams),
99+
new ChromeAdapter(
100+
window.LanguageModel as LanguageModel,
101+
hybridParams.mode,
102+
hybridParams.onDeviceParams
103+
),
99104
requestOptions
100105
);
101106
}

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

+310
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,310 @@
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 { expect, use } from 'chai';
19+
import sinonChai from 'sinon-chai';
20+
import chaiAsPromised from 'chai-as-promised';
21+
import { ChromeAdapter } from './chrome-adapter';
22+
import {
23+
Availability,
24+
LanguageModel,
25+
LanguageModelCreateOptions
26+
} from '../types/language-model';
27+
import { stub } from 'sinon';
28+
import { GenerateContentRequest } from '../types';
29+
30+
use(sinonChai);
31+
use(chaiAsPromised);
32+
33+
describe('ChromeAdapter', () => {
34+
describe('isAvailable', () => {
35+
it('returns false if mode is only cloud', async () => {
36+
const adapter = new ChromeAdapter(undefined, 'only_in_cloud');
37+
expect(
38+
await adapter.isAvailable({
39+
contents: []
40+
})
41+
).to.be.false;
42+
});
43+
it('returns false if AI API is undefined', async () => {
44+
const adapter = new ChromeAdapter(undefined, 'prefer_on_device');
45+
expect(
46+
await adapter.isAvailable({
47+
contents: []
48+
})
49+
).to.be.false;
50+
});
51+
it('returns false if LanguageModel API is undefined', async () => {
52+
const adapter = new ChromeAdapter(
53+
{} as LanguageModel,
54+
'prefer_on_device'
55+
);
56+
expect(
57+
await adapter.isAvailable({
58+
contents: []
59+
})
60+
).to.be.false;
61+
});
62+
it('returns false if request contents empty', async () => {
63+
const adapter = new ChromeAdapter(
64+
{} as LanguageModel,
65+
'prefer_on_device'
66+
);
67+
expect(
68+
await adapter.isAvailable({
69+
contents: []
70+
})
71+
).to.be.false;
72+
});
73+
it('returns false if request content has function role', async () => {
74+
const adapter = new ChromeAdapter(
75+
{} as LanguageModel,
76+
'prefer_on_device'
77+
);
78+
expect(
79+
await adapter.isAvailable({
80+
contents: [
81+
{
82+
role: 'function',
83+
parts: []
84+
}
85+
]
86+
})
87+
).to.be.false;
88+
});
89+
it('returns false if request content has multiple parts', async () => {
90+
const adapter = new ChromeAdapter(
91+
{} as LanguageModel,
92+
'prefer_on_device'
93+
);
94+
expect(
95+
await adapter.isAvailable({
96+
contents: [
97+
{
98+
role: 'user',
99+
parts: [{ text: 'a' }, { text: 'b' }]
100+
}
101+
]
102+
})
103+
).to.be.false;
104+
});
105+
it('returns false if request content has non-text part', async () => {
106+
const adapter = new ChromeAdapter(
107+
{} as LanguageModel,
108+
'prefer_on_device'
109+
);
110+
expect(
111+
await adapter.isAvailable({
112+
contents: [
113+
{
114+
role: 'user',
115+
parts: [{ inlineData: { mimeType: 'a', data: 'b' } }]
116+
}
117+
]
118+
})
119+
).to.be.false;
120+
});
121+
it('returns false if request system instruction has function role', async () => {
122+
const adapter = new ChromeAdapter(
123+
{} as LanguageModel,
124+
'prefer_on_device'
125+
);
126+
expect(
127+
await adapter.isAvailable({
128+
contents: [],
129+
systemInstruction: {
130+
role: 'function',
131+
parts: []
132+
}
133+
})
134+
).to.be.false;
135+
});
136+
it('returns false if request system instruction has multiple parts', async () => {
137+
const adapter = new ChromeAdapter(
138+
{} as LanguageModel,
139+
'prefer_on_device'
140+
);
141+
expect(
142+
await adapter.isAvailable({
143+
contents: [],
144+
systemInstruction: {
145+
role: 'function',
146+
parts: [{ text: 'a' }, { text: 'b' }]
147+
}
148+
})
149+
).to.be.false;
150+
});
151+
it('returns false if request system instruction has non-text part', async () => {
152+
const adapter = new ChromeAdapter(
153+
{} as LanguageModel,
154+
'prefer_on_device'
155+
);
156+
expect(
157+
await adapter.isAvailable({
158+
contents: [],
159+
systemInstruction: {
160+
role: 'function',
161+
parts: [{ inlineData: { mimeType: 'a', data: 'b' } }]
162+
}
163+
})
164+
).to.be.false;
165+
});
166+
it('returns true if model is readily available', async () => {
167+
const languageModelProvider = {
168+
availability: () => Promise.resolve(Availability.available)
169+
} as LanguageModel;
170+
const adapter = new ChromeAdapter(
171+
languageModelProvider,
172+
'prefer_on_device'
173+
);
174+
expect(
175+
await adapter.isAvailable({
176+
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
177+
})
178+
).to.be.true;
179+
});
180+
it('returns false and triggers download when model is available after download', async () => {
181+
const languageModelProvider = {
182+
availability: () => Promise.resolve(Availability.downloadable),
183+
create: () => Promise.resolve({})
184+
} as LanguageModel;
185+
const createStub = stub(languageModelProvider, 'create').resolves(
186+
{} as LanguageModel
187+
);
188+
const onDeviceParams = {} as LanguageModelCreateOptions;
189+
const adapter = new ChromeAdapter(
190+
languageModelProvider,
191+
'prefer_on_device',
192+
onDeviceParams
193+
);
194+
expect(
195+
await adapter.isAvailable({
196+
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
197+
})
198+
).to.be.false;
199+
expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
200+
});
201+
it('avoids redundant downloads', async () => {
202+
const languageModelProvider = {
203+
availability: () => Promise.resolve(Availability.downloadable),
204+
create: () => Promise.resolve({})
205+
} as LanguageModel;
206+
const downloadPromise = new Promise<LanguageModel>(() => {
207+
/* never resolves */
208+
});
209+
const createStub = stub(languageModelProvider, 'create').returns(
210+
downloadPromise
211+
);
212+
const adapter = new ChromeAdapter(languageModelProvider);
213+
await adapter.isAvailable({
214+
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
215+
});
216+
await adapter.isAvailable({
217+
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
218+
});
219+
expect(createStub).to.have.been.calledOnce;
220+
});
221+
it('clears state when download completes', async () => {
222+
const languageModelProvider = {
223+
availability: () => Promise.resolve(Availability.downloadable),
224+
create: () => Promise.resolve({})
225+
} as LanguageModel;
226+
let resolveDownload;
227+
const downloadPromise = new Promise<LanguageModel>(resolveCallback => {
228+
resolveDownload = resolveCallback;
229+
});
230+
const createStub = stub(languageModelProvider, 'create').returns(
231+
downloadPromise
232+
);
233+
const adapter = new ChromeAdapter(languageModelProvider);
234+
await adapter.isAvailable({
235+
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
236+
});
237+
resolveDownload!();
238+
await adapter.isAvailable({
239+
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
240+
});
241+
expect(createStub).to.have.been.calledTwice;
242+
});
243+
it('returns false when model is never available', async () => {
244+
const languageModelProvider = {
245+
availability: () => Promise.resolve(Availability.unavailable),
246+
create: () => Promise.resolve({})
247+
} as LanguageModel;
248+
const adapter = new ChromeAdapter(
249+
languageModelProvider,
250+
'prefer_on_device'
251+
);
252+
expect(
253+
await adapter.isAvailable({
254+
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
255+
})
256+
).to.be.false;
257+
});
258+
});
259+
describe('generateContentOnDevice', () => {
260+
it('generates content', async () => {
261+
const languageModelProvider = {
262+
create: () => Promise.resolve({})
263+
} as LanguageModel;
264+
const languageModel = {
265+
prompt: i => Promise.resolve(i)
266+
} as LanguageModel;
267+
const createStub = stub(languageModelProvider, 'create').resolves(
268+
languageModel
269+
);
270+
const promptOutput = 'hi';
271+
const promptStub = stub(languageModel, 'prompt').resolves(promptOutput);
272+
const onDeviceParams = {
273+
systemPrompt: 'be yourself'
274+
} as LanguageModelCreateOptions;
275+
const adapter = new ChromeAdapter(
276+
languageModelProvider,
277+
'prefer_on_device',
278+
onDeviceParams
279+
);
280+
const request = {
281+
contents: [{ role: 'user', parts: [{ text: 'anything' }] }]
282+
} as GenerateContentRequest;
283+
const response = await adapter.generateContentOnDevice(request);
284+
// Asserts initialization params are proxied.
285+
expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
286+
// Asserts Vertex input type is mapped to Chrome type.
287+
expect(promptStub).to.have.been.calledOnceWith([
288+
{
289+
role: request.contents[0].role,
290+
content: [
291+
{
292+
type: 'text',
293+
content: request.contents[0].parts[0].text
294+
}
295+
]
296+
}
297+
]);
298+
// Asserts expected output.
299+
expect(await response.json()).to.deep.equal({
300+
candidates: [
301+
{
302+
content: {
303+
parts: [{ text: promptOutput }]
304+
}
305+
}
306+
]
307+
});
308+
});
309+
});
310+
});

0 commit comments

Comments
 (0)