|
16 | 16 | */
|
17 | 17 |
|
18 | 18 | import { expect } from 'chai';
|
19 |
| -import { Content, GenerationConfig, HarmBlockMethod, HarmBlockThreshold, HarmCategory, Modality, SafetySetting, getAI, getGenerativeModel, getVertexAI } from '../src'; |
20 | 19 | 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, |
21 | 39 | testConfigs
|
22 | 40 | } from './constants';
|
| 41 | +import { FIREBASE_CONFIG } from './firebase-config'; |
| 42 | + |
23 | 43 |
|
24 | 44 | describe('Count Tokens', () => {
|
25 | 45 | testConfigs.forEach(testConfig => {
|
@@ -77,30 +97,86 @@ describe('Count Tokens', () => {
|
77 | 97 | expect(response.promptTokensDetails![0].modality).to.equal(Modality.TEXT);
|
78 | 98 | expect(response.promptTokensDetails![0].tokenCount).to.equal(6);
|
79 | 99 | });
|
| 100 | + |
80 | 101 | 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 | + }); |
81 | 118 |
|
82 |
| - }) |
83 | 119 | 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 | + }); |
84 | 136 |
|
85 |
| - }) |
86 | 137 | 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 } }; |
87 | 142 |
|
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); |
90 | 152 |
|
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); |
93 | 156 |
|
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 | + }); |
96 | 161 |
|
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 | + }); |
104 | 180 | });
|
105 | 181 | })
|
106 | 182 | });
|
0 commit comments