Skip to content

Commit 7d9a71a

Browse files
erikeldridgeerikeldridge
committed
Test request-based availability checks
1 parent dd9ac27 commit 7d9a71a

File tree

2 files changed

+157
-32
lines changed

2 files changed

+157
-32
lines changed

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

Lines changed: 136 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,29 +21,155 @@ import chaiAsPromised from 'chai-as-promised';
2121
import { ChromeAdapter } from './chrome-adapter';
2222
import { Availability, LanguageModel } from '../types/language-model';
2323
import { stub } from 'sinon';
24+
import * as util from '@firebase/util';
2425

2526
use(sinonChai);
2627
use(chaiAsPromised);
2728

2829
describe('ChromeAdapter', () => {
29-
describe('isOnDeviceRequest', () => {
30-
it('returns true for simple text part', async () => {
30+
describe('isAvailable', () => {
31+
it('returns false if mode is only cloud', async () => {
32+
const adapter = new ChromeAdapter(undefined, 'only_in_cloud');
3133
expect(
32-
ChromeAdapter._isOnDeviceRequest({
33-
contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
34+
await adapter.isAvailable({
35+
contents: []
3436
})
35-
).to.be.true;
37+
).to.be.false;
3638
});
37-
it('returns false if contents empty', async () => {
39+
it('returns false if browser is not Chrome', async () => {
40+
const chromeStub = stub(util, 'isChrome').returns(false);
41+
const adapter = new ChromeAdapter(undefined, 'prefer_on_device');
3842
expect(
39-
ChromeAdapter._isOnDeviceRequest({
43+
await adapter.isAvailable({
4044
contents: []
4145
})
4246
).to.be.false;
47+
chromeStub.restore();
4348
});
44-
});
45-
describe('isAvailable', () => {
46-
it('returns true if a model is available', async () => {
49+
it('returns false if AI API is undefined', async () => {
50+
const adapter = new ChromeAdapter(undefined, 'prefer_on_device');
51+
expect(
52+
await adapter.isAvailable({
53+
contents: []
54+
})
55+
).to.be.false;
56+
});
57+
it('returns false if LanguageModel API is undefined', async () => {
58+
const adapter = new ChromeAdapter(
59+
{} as LanguageModel,
60+
'prefer_on_device'
61+
);
62+
expect(
63+
await adapter.isAvailable({
64+
contents: []
65+
})
66+
).to.be.false;
67+
});
68+
it('returns false if request contents empty', async () => {
69+
const adapter = new ChromeAdapter(
70+
{} as LanguageModel,
71+
'prefer_on_device'
72+
);
73+
expect(
74+
await adapter.isAvailable({
75+
contents: []
76+
})
77+
).to.be.false;
78+
});
79+
it('returns false if request content has function role', async () => {
80+
const adapter = new ChromeAdapter(
81+
{} as LanguageModel,
82+
'prefer_on_device'
83+
);
84+
expect(
85+
await adapter.isAvailable({
86+
contents: [
87+
{
88+
role: 'function',
89+
parts: []
90+
}
91+
]
92+
})
93+
).to.be.false;
94+
});
95+
it('returns false if request content has multiple parts', async () => {
96+
const adapter = new ChromeAdapter(
97+
{} as LanguageModel,
98+
'prefer_on_device'
99+
);
100+
expect(
101+
await adapter.isAvailable({
102+
contents: [
103+
{
104+
role: 'user',
105+
parts: [{ text: 'a' }, { text: 'b' }]
106+
}
107+
]
108+
})
109+
).to.be.false;
110+
});
111+
it('returns false if request content has non-text part', async () => {
112+
const adapter = new ChromeAdapter(
113+
{} as LanguageModel,
114+
'prefer_on_device'
115+
);
116+
expect(
117+
await adapter.isAvailable({
118+
contents: [
119+
{
120+
role: 'user',
121+
parts: [{ inlineData: { mimeType: 'a', data: 'b' } }]
122+
}
123+
]
124+
})
125+
).to.be.false;
126+
});
127+
it('returns false if request system instruction has function role', async () => {
128+
const adapter = new ChromeAdapter(
129+
{} as LanguageModel,
130+
'prefer_on_device'
131+
);
132+
expect(
133+
await adapter.isAvailable({
134+
contents: [],
135+
systemInstruction: {
136+
role: 'function',
137+
parts: []
138+
}
139+
})
140+
).to.be.false;
141+
});
142+
it('returns false if request system instruction has multiple parts', async () => {
143+
const adapter = new ChromeAdapter(
144+
{} as LanguageModel,
145+
'prefer_on_device'
146+
);
147+
expect(
148+
await adapter.isAvailable({
149+
contents: [],
150+
systemInstruction: {
151+
role: 'function',
152+
parts: [{ text: 'a' }, { text: 'b' }]
153+
}
154+
})
155+
).to.be.false;
156+
});
157+
it('returns false if request system instruction has non-text part', async () => {
158+
const adapter = new ChromeAdapter(
159+
{} as LanguageModel,
160+
'prefer_on_device'
161+
);
162+
expect(
163+
await adapter.isAvailable({
164+
contents: [],
165+
systemInstruction: {
166+
role: 'function',
167+
parts: [{ inlineData: { mimeType: 'a', data: 'b' } }]
168+
}
169+
})
170+
).to.be.false;
171+
});
172+
it('returns true if model is readily available', async () => {
47173
const languageModelProvider = {
48174
availability: () => Promise.resolve(Availability.available)
49175
} as LanguageModel;

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

Lines changed: 21 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ export class ChromeAdapter {
6767
return false;
6868
}
6969
// Returns false if the request can't be run on-device.
70-
if (!ChromeAdapter._isOnDeviceRequest(request)) {
70+
if (!ChromeAdapter.isOnDeviceRequest(request)) {
7171
return false;
7272
}
7373
const availability = await this.languageModelProvider.availability();
@@ -109,27 +109,7 @@ export class ChromeAdapter {
109109
})
110110
} as Response;
111111
}
112-
// Visible for testing
113-
static _isOnDeviceRequest(request: GenerateContentRequest): boolean {
114-
if (request.systemInstruction) {
115-
const systemContent = request.systemInstruction as Content;
116-
// Returns false if the role can't be represented on-device.
117-
if (systemContent.role && systemContent.role === 'function') {
118-
return false;
119-
}
120-
121-
// Returns false if the system prompt is multi-part.
122-
if (systemContent.parts && systemContent.parts.length > 1) {
123-
return false;
124-
}
125-
126-
// Returns false if the system prompt isn't text.
127-
const systemText = request.systemInstruction as TextPart;
128-
if (!systemText.text) {
129-
return false;
130-
}
131-
}
132-
112+
private static isOnDeviceRequest(request: GenerateContentRequest): boolean {
133113
// Returns false if the prompt is empty.
134114
if (request.contents.length === 0) {
135115
return false;
@@ -150,6 +130,25 @@ export class ChromeAdapter {
150130
}
151131
}
152132

133+
if (request.systemInstruction) {
134+
const systemContent = request.systemInstruction as Content;
135+
// Returns false if the role can't be represented on-device.
136+
if (systemContent.role && systemContent.role === 'function') {
137+
return false;
138+
}
139+
140+
// Returns false if the system prompt is multi-part.
141+
if (systemContent.parts && systemContent.parts.length > 1) {
142+
return false;
143+
}
144+
145+
// Returns false if the system prompt isn't text.
146+
const systemText = request.systemInstruction as TextPart;
147+
if (!systemText.text) {
148+
return false;
149+
}
150+
}
151+
153152
return true;
154153
}
155154
private download(): void {

0 commit comments

Comments
 (0)