Skip to content

Commit d78fc29

Browse files
dlarocquedlarocque
committed
count tokens test cleanup
1 parent c769317 commit d78fc29

File tree

3 files changed

+62
-36
lines changed

3 files changed

+62
-36
lines changed

‎packages/ai/integration/constants.ts

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

1818
import { initializeApp } from '@firebase/app';
19-
import { AI, Backend, BackendType, GoogleAIBackend, VertexAIBackend, getAI } from '../src';
19+
import {
20+
AI,
21+
Backend,
22+
BackendType,
23+
GoogleAIBackend,
24+
VertexAIBackend,
25+
getAI
26+
} from '../src';
2027
import { FIREBASE_CONFIG } from './firebase-config';
2128

2229
const app = initializeApp(FIREBASE_CONFIG);
@@ -51,7 +58,7 @@ const modelNames: readonly string[] = [
5158
];
5259

5360
/**
54-
* Array of test configurations that is iterated over to get full coverage
61+
* Array of test configurations that is iterated over to get full coverage
5562
* of backends and models. Contains all combinations of backends and models.
5663
*/
5764
export const testConfigs: readonly TestConfig[] = backends.flatMap(backend => {

‎packages/ai/integration/count-tokens.test.ts

Lines changed: 51 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -88,17 +88,18 @@ describe('Count Tokens', () => {
8888

8989
const response = await model.countTokens('Why is the sky blue?');
9090

91-
expect(response.promptTokensDetails).to.not.be.null;
91+
expect(response.promptTokensDetails).to.exist;
9292
expect(response.promptTokensDetails!.length).to.equal(1);
9393
expect(response.promptTokensDetails![0].modality).to.equal(
9494
Modality.TEXT
9595
);
96-
9796
if (testConfig.ai.backend.backendType === BackendType.GOOGLE_AI) {
9897
expect(response.totalTokens).to.equal(7);
9998
expect(response.totalBillableCharacters).to.be.undefined;
10099
expect(response.promptTokensDetails![0].tokenCount).to.equal(7);
101-
} else if (testConfig.ai.backend.backendType === BackendType.VERTEX_AI) {
100+
} else if (
101+
testConfig.ai.backend.backendType === BackendType.VERTEX_AI
102+
) {
102103
expect(response.totalTokens).to.equal(6);
103104
expect(response.totalBillableCharacters).to.equal(16);
104105
expect(response.promptTokensDetails![0].tokenCount).to.equal(6);
@@ -119,7 +120,17 @@ describe('Count Tokens', () => {
119120

120121
if (testConfig.ai.backend.backendType === BackendType.GOOGLE_AI) {
121122
const expectedImageTokens = 259;
122-
123+
expect(response.totalTokens).to.equal(expectedImageTokens);
124+
expect(response.totalBillableCharacters).to.be.undefined; // Incorrect behavior
125+
expect(response.promptTokensDetails!.length).to.equal(2);
126+
expect(response.promptTokensDetails![0]).to.deep.equal({
127+
modality: Modality.TEXT, // Note: 1 unexpected text token observed for Google AI with image-only input.
128+
tokenCount: 1
129+
});
130+
expect(response.promptTokensDetails![1]).to.deep.equal({
131+
modality: Modality.IMAGE,
132+
tokenCount: 258
133+
});
123134
} else if (testConfig.ai.backend.backendType === BackendType.VERTEX_AI) {
124135
const expectedImageTokens = 258;
125136
expect(response.totalTokens).to.equal(expectedImageTokens);
@@ -129,9 +140,8 @@ describe('Count Tokens', () => {
129140
expect(
130141
response.promptTokensDetails!.length,
131142
).to.equal(1);
132-
expect(
133-
response.promptTokensDetails![0].modality,
134-
).to.equal(Modality.IMAGE);
143+
// Note: No text tokens are present for Vertex AI with image-only input.
144+
expect(response.promptTokensDetails![0]).to.deep.equal({ modality: Modality.IMAGE, tokenCount: 258 })
135145
expect(response.promptTokensDetails![0].tokenCount).to.equal(expectedImageTokens);
136146
}
137147
});
@@ -149,6 +159,7 @@ describe('Count Tokens', () => {
149159

150160
const response = await model.countTokens([audioPart]);
151161

162+
expect(response.promptTokensDetails).to.exist;
152163
const textDetails = response.promptTokensDetails!.find(
153164
d => d.modality === Modality.TEXT
154165
);
@@ -158,20 +169,24 @@ describe('Count Tokens', () => {
158169

159170
if (testConfig.ai.backend.backendType === BackendType.GOOGLE_AI) {
160171
expect(response.totalTokens).to.equal(6);
161-
expect(
162-
response.promptTokensDetails!.length,
163-
).to.equal(2);
164-
expect(textDetails).to.deep.equal({ modality: Modality.TEXT, tokenCount: 1 })
165-
expect(audioDetails).to.deep.equal({ modality: Modality.AUDIO, tokenCount: 5 })
166-
} else if (testConfig.ai.backend.backendType === BackendType.VERTEX_AI) {
172+
expect(response.promptTokensDetails!.length).to.equal(2);
173+
expect(textDetails).to.deep.equal({
174+
modality: Modality.TEXT,
175+
tokenCount: 1
176+
});
177+
expect(audioDetails).to.deep.equal({
178+
modality: Modality.AUDIO,
179+
tokenCount: 5
180+
});
181+
} else if (
182+
testConfig.ai.backend.backendType === BackendType.VERTEX_AI
183+
) {
167184
expect(response.totalTokens).to.be.undefined;
168-
expect(response.promptTokensDetails!.length).to.equal(1); // For some reason we don't get text
169-
expect(audioDetails).to.deep.equal({ modality: Modality.AUDIO }); // For some reason there are no tokens
185+
expect(response.promptTokensDetails!.length).to.equal(1); // Note: Text modality details absent for Vertex AI with audio-only input.
186+
expect(audioDetails).to.deep.equal({ modality: Modality.AUDIO }); // Note: Audio tokenCount is undefined for Vertex AI with audio-only input.
170187
}
171188

172-
expect(
173-
response.totalBillableCharacters,
174-
).to.be.undefined; // Incorrect behavior
189+
expect(response.totalBillableCharacters).to.be.undefined; // Incorrect behavior
175190
});
176191

177192
it('text, image, and audio input', async () => {
@@ -193,12 +208,19 @@ describe('Count Tokens', () => {
193208
const textDetails = response.promptTokensDetails!.find(
194209
d => d.modality === Modality.TEXT
195210
);
196-
const visionDetails = response.promptTokensDetails!.find(
211+
const imageDetails = response.promptTokensDetails!.find(
197212
d => d.modality === Modality.IMAGE
198213
);
199214
const audioDetails = response.promptTokensDetails!.find(
200215
d => d.modality === Modality.AUDIO
201216
);
217+
expect(response.promptTokensDetails).to.exist;
218+
expect(response.promptTokensDetails!.length).to.equal(3);
219+
220+
expect(imageDetails).to.deep.equal({
221+
modality: Modality.IMAGE,
222+
tokenCount: 258
223+
});
202224

203225
if (testConfig.ai.backend.backendType === BackendType.GOOGLE_AI) {
204226
expect(response.totalTokens).to.equal(267);
@@ -207,25 +229,22 @@ describe('Count Tokens', () => {
207229
modality: Modality.TEXT,
208230
tokenCount: 4
209231
});
210-
expect(audioDetails).to.deep.equal({ modality: Modality.AUDIO, tokenCount: 5 }); // Incorrect behavior because there's no tokenCount
211-
} else if (testConfig.ai.backend.backendType === BackendType.VERTEX_AI) {
232+
expect(audioDetails).to.deep.equal({
233+
modality: Modality.AUDIO,
234+
tokenCount: 5
235+
});
236+
} else if (
237+
testConfig.ai.backend.backendType === BackendType.VERTEX_AI
238+
) {
212239
expect(response.totalTokens).to.equal(261);
213240
expect(textDetails).to.deep.equal({
214241
modality: Modality.TEXT,
215242
tokenCount: 3
216243
});
217-
expect(
218-
response.totalBillableCharacters,
219-
).to.equal('Describe these:'.length - 1); // For some reason it's the length-1
244+
const expectedText = 'Describe these:';
245+
expect(response.totalBillableCharacters).to.equal(expectedText.length - 1); // Note: BillableCharacters observed as (text length - 1) for Vertex AI.
220246
expect(audioDetails).to.deep.equal({ modality: Modality.AUDIO }); // Incorrect behavior because there's no tokenCount
221247
}
222-
223-
expect(response.promptTokensDetails!.length).to.equal(3);
224-
225-
expect(visionDetails).to.deep.equal({
226-
modality: Modality.IMAGE,
227-
tokenCount: 258
228-
});
229248
});
230249

231250
it('public storage reference', async () => {
@@ -248,7 +267,7 @@ describe('Count Tokens', () => {
248267
const expectedFileTokens = 258;
249268
expect(response.totalTokens).to.equal(expectedFileTokens);
250269
expect(response.totalBillableCharacters).to.be.undefined;
251-
expect(response.promptTokensDetails).to.not.be.null;
270+
expect(response.promptTokensDetails).to.exist;
252271
expect(response.promptTokensDetails!.length).to.equal(1);
253272
expect(response.promptTokensDetails![0].modality).to.equal(
254273
Modality.IMAGE

‎packages/ai/integration/generate-content.test.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,11 @@ describe('Generate Content', () => {
4545
const safetySettings: SafetySetting[] = [
4646
{
4747
category: HarmCategory.HARM_CATEGORY_HARASSMENT,
48-
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
48+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
4949
},
5050
{
5151
category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
52-
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
52+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
5353
},
5454
{
5555
category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,

0 commit comments

Comments
 (0)