Skip to content

Commit 472c452

Browse files
committed
count tokens tests
1 parent abe830c commit 472c452

File tree

2 files changed

+102
-23
lines changed

2 files changed

+102
-23
lines changed

packages/ai/integration/constants.ts

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -27,19 +27,17 @@ import { FIREBASE_CONFIG } from './firebase-config';
2727

2828
const app = initializeApp(FIREBASE_CONFIG);
2929

30-
export type ModelName = 'gemini-2.0-flash' | 'gemini-2.0-flash-exp';
31-
3230
/**
3331
* Test config that all tests will be ran against.
3432
*/
3533
export type TestConfig = Readonly<{
3634
ai: AI;
37-
model: ModelName;
35+
model: string;
3836
/** This will be used to output the test config at runtime */
3937
toString: () => string;
4038
}>;
4139

42-
function formatConfigAsString(config: { ai: AI; model: ModelName }): string {
40+
function formatConfigAsString(config: { ai: AI; model: string }): string {
4341
return `${backendNames.get(config.ai.backend.backendType)} ${config.model}`;
4442
}
4543

@@ -53,9 +51,9 @@ const backendNames: Map<BackendType, string> = new Map([
5351
[BackendType.VERTEX_AI, 'Vertex AI']
5452
]);
5553

56-
const modelNames: ReadonlyArray<ModelName> = [
54+
const modelNames: ReadonlyArray<string> = [
5755
'gemini-2.0-flash',
58-
'gemini-2.0-flash-exp'
56+
// 'gemini-2.0-flash-exp'
5957
];
6058

6159
export const testConfigs: ReadonlyArray<TestConfig> = backends.flatMap(backend => {
@@ -67,4 +65,9 @@ export const testConfigs: ReadonlyArray<TestConfig> = backends.flatMap(backend =
6765
toString: () => formatConfigAsString({ ai, model: modelName })
6866
}
6967
})
70-
})
68+
})
69+
70+
export const TINY_IMG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABAQMAAAAl21bKAAAAA1BMVEUAAACnej3aAAAAAXRSTlMAQObYZgAAAApJREFUCNdjYAAAAAIAAeIhvDMAAAAASUVORK5CYII=';
71+
export const IMAGE_MIME_TYPE = 'image/png';
72+
export const TINY_MP3_BASE64 = 'SUQzBAAAAAAAIlRTU0UAAAAOAAADTGF2ZjYxLjcuMTAwAAAAAAAAAAAAAAD/+0DAAAAAAAAAAAAAAAAAAAAAAABJbmZvAAAADwAAAAUAAAK+AGhoaGhoaGhoaGhoaGhoaGhoaGiOjo6Ojo6Ojo6Ojo6Ojo6Ojo6OjrS0tLS0tLS0tLS0tLS0tLS0tLS02tra2tra2tra2tra2tra2tra2tr//////////////////////////wAAAABMYXZjNjEuMTkAAAAAAAAAAAAAAAAkAwYAAAAAAAACvhC6DYoAAAAAAP/7EMQAA8AAAaQAAAAgAAA0gAAABExBTUUzLjEwMFVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVV//sQxCmDwAABpAAAACAAADSAAAAEVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVX/+xDEUwPAAAGkAAAAIAAANIAAAARVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVf/7EMR8g8AAAaQAAAAgAAA0gAAABFVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVV//sQxKYDwAABpAAAACAAADSAAAAEVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVU=';
73+
export const AUDIO_MIME_TYPE = 'audio/mpeg';

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

Lines changed: 92 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,30 @@
1616
*/
1717

1818
import { expect } from 'chai';
19-
import { Content, GenerationConfig, HarmBlockMethod, HarmBlockThreshold, HarmCategory, Modality, SafetySetting, getAI, getGenerativeModel, getVertexAI } from '../src';
2019
import {
20+
Content,
21+
GenerationConfig,
22+
HarmBlockMethod,
23+
HarmBlockThreshold,
24+
HarmCategory,
25+
Modality,
26+
SafetySetting,
27+
getGenerativeModel,
28+
Part,
29+
CountTokensRequest,
30+
Schema,
31+
InlineDataPart,
32+
FileDataPart
33+
} from '../src';
34+
import {
35+
AUDIO_MIME_TYPE,
36+
IMAGE_MIME_TYPE,
37+
TINY_IMG_BASE64,
38+
TINY_MP3_BASE64,
2139
testConfigs
2240
} from './constants';
41+
import { FIREBASE_CONFIG } from './firebase-config';
42+
2343

2444
describe('Count Tokens', () => {
2545
testConfigs.forEach(testConfig => {
@@ -77,30 +97,86 @@ describe('Count Tokens', () => {
7797
expect(response.promptTokensDetails![0].modality).to.equal(Modality.TEXT);
7898
expect(response.promptTokensDetails![0].tokenCount).to.equal(6);
7999
});
100+
80101
it('image input', async () => {
102+
const model = getGenerativeModel(testConfig.ai, { model: testConfig.model });
103+
const imagePart: Part = {
104+
inlineData: {
105+
mimeType: IMAGE_MIME_TYPE,
106+
data: TINY_IMG_BASE64
107+
}
108+
};
109+
const response = await model.countTokens([imagePart]);
110+
111+
const expectedImageTokens = 258;
112+
expect(response.totalTokens, 'totalTokens should have correct token count').to.equal(expectedImageTokens);
113+
expect(response.totalBillableCharacters, 'totalBillableCharacters should be undefined').to.be.undefined; // Incorrect behavior
114+
expect(response.promptTokensDetails!.length, 'promptTokensDetails should have one entry').to.equal(1);
115+
expect(response.promptTokensDetails![0].modality, 'modality should be IMAGE').to.equal(Modality.IMAGE);
116+
expect(response.promptTokensDetails![0].tokenCount, 'promptTokenDetails tokenCount should be correct').to.equal(expectedImageTokens);
117+
});
81118

82-
})
83119
it('audio input', async () => {
120+
const model = getGenerativeModel(testConfig.ai, { model: testConfig.model });
121+
const audioPart: InlineDataPart = {
122+
inlineData: {
123+
mimeType: AUDIO_MIME_TYPE,
124+
data: TINY_MP3_BASE64
125+
}
126+
};
127+
128+
const response = await model.countTokens([audioPart]);
129+
// This may be different on Google AI
130+
expect(response.totalTokens, 'totalTokens is expected to be undefined').to.be.undefined;
131+
expect(response.totalBillableCharacters, 'totalBillableCharacters should be undefined').to.be.undefined; // Incorrect behavior
132+
expect(response.promptTokensDetails!.length, 'promptTokensDetails should have one entry').to.equal(1);
133+
expect(response.promptTokensDetails![0].modality, 'modality should be AUDIO').to.equal(Modality.AUDIO);
134+
expect(response.promptTokensDetails![0].tokenCount, 'promptTokenDetails tokenCount is expected to be undefined').to.be.undefined;
135+
});
84136

85-
})
86137
it('text, image, and audio input', async () => {
138+
const model = getGenerativeModel(testConfig.ai, { model: testConfig.model });
139+
const textPart: Part = { text: 'Describe these:' };
140+
const imagePart: Part = { inlineData: { mimeType: IMAGE_MIME_TYPE, data: TINY_IMG_BASE64 } };
141+
const audioPart: Part = { inlineData: { mimeType: AUDIO_MIME_TYPE, data: TINY_MP3_BASE64 } };
87142

88-
})
89-
it('public storage reference', async () => {
143+
const request: CountTokensRequest = {
144+
contents: [{ role: 'user', parts: [textPart, imagePart, audioPart] }]
145+
};
146+
const response = await model.countTokens(request);
147+
148+
expect(response.totalTokens, 'totalTokens should have correct token count').to.equal(261);
149+
expect(response.totalBillableCharacters, 'totalBillableCharacters should have correct count').to.equal('Describe these:'.length - 1); // For some reason it's the length-1
150+
151+
expect(response.promptTokensDetails!.length, 'promptTokensDetails should have three entries').to.equal(3);
90152

91-
})
92-
it('private storage reference', async () => {
153+
const textDetails = response.promptTokensDetails!.find(d => d.modality === Modality.TEXT);
154+
const visionDetails = response.promptTokensDetails!.find(d => d.modality === Modality.IMAGE);
155+
const audioDetails = response.promptTokensDetails!.find(d => d.modality === Modality.AUDIO);
93156

94-
})
95-
it('schema', async () => {
157+
expect(textDetails).to.deep.equal({ modality: Modality.TEXT, tokenCount: 3 });
158+
expect(visionDetails).to.deep.equal({ modality: Modality.IMAGE, tokenCount: 258 });
159+
expect(audioDetails).to.deep.equal({ modality: Modality.AUDIO }); // Incorrect behavior because there's no tokenCount
160+
});
96161

97-
})
98-
// TODO (dlarocque): Test countTokens() with the following:
99-
// - inline data
100-
// - public storage reference
101-
// - private storage reference (testing auth integration)
102-
// - count tokens
103-
// - JSON schema
162+
it('public storage reference', async () => {
163+
const model = getGenerativeModel(testConfig.ai, { model: testConfig.model });
164+
const filePart: FileDataPart = {
165+
fileData: {
166+
mimeType: IMAGE_MIME_TYPE,
167+
fileUri: `gs://${FIREBASE_CONFIG.storageBucket}/images/tree.png`
168+
}
169+
};
170+
const response = await model.countTokens([filePart]);
171+
172+
const expectedFileTokens = 258;
173+
expect(response.totalTokens, 'totalTokens should have correct token count').to.equal(expectedFileTokens);
174+
expect(response.totalBillableCharacters, 'totalBillableCharacters should be undefined').to.be.undefined;
175+
expect(response.promptTokensDetails).to.not.be.null;
176+
expect(response.promptTokensDetails!.length).to.equal(1);
177+
expect(response.promptTokensDetails![0].modality).to.equal(Modality.IMAGE);
178+
expect(response.promptTokensDetails![0].tokenCount).to.equal(expectedFileTokens);
179+
});
104180
});
105181
})
106182
});

0 commit comments

Comments
 (0)