Skip to content

Commit 9f908d2

Browse files
dlarocquedlarocque
committed
Add test configs to run against each test
1 parent 11a8b05 commit 9f908d2

File tree

8 files changed

+279
-222
lines changed

8 files changed

+279
-222
lines changed

‎packages/ai/integration/constants.ts

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
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 { initializeApp } from '@firebase/app';
19+
import {
20+
AI,
21+
Backend,
22+
BackendType,
23+
VertexAIBackend,
24+
getAI
25+
} from '../src';
26+
import { FIREBASE_CONFIG } from './firebase-config';
27+
28+
const app = initializeApp(FIREBASE_CONFIG);
29+
30+
export type ModelName = 'gemini-2.0-flash' | 'gemini-2.0-flash-exp';
31+
32+
/**
33+
* Test config that all tests will be ran against.
34+
*/
35+
export type TestConfig = Readonly<{
36+
ai: AI;
37+
model: ModelName;
38+
/** This will be used to output the test config at runtime */
39+
toString: () => string;
40+
}>;
41+
42+
function formatConfigAsString(config: { ai: AI; model: ModelName }): string {
43+
return `${backendNames.get(config.ai.backend.backendType)} ${config.model}`;
44+
}
45+
46+
const backends: ReadonlyArray<Backend> = [
47+
// new GoogleAIBackend(), TODO: activate once live
48+
new VertexAIBackend()
49+
];
50+
51+
const backendNames: Map<BackendType, string> = new Map([
52+
[BackendType.GOOGLE_AI, 'Google AI'],
53+
[BackendType.VERTEX_AI, 'Vertex AI']
54+
]);
55+
56+
const modelNames: ReadonlyArray<ModelName> = [
57+
'gemini-2.0-flash',
58+
'gemini-2.0-flash-exp'
59+
];
60+
61+
export const testConfigs: ReadonlyArray<TestConfig> = backends.flatMap(backend => {
62+
return modelNames.map(modelName => {
63+
const ai = getAI(app, { backend });
64+
return {
65+
ai: getAI(app, { backend }),
66+
model: modelName,
67+
toString: () => formatConfigAsString({ ai, model: modelName })
68+
}
69+
})
70+
})

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

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
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 } from 'chai';
19+
import { Content, GenerationConfig, HarmBlockMethod, HarmBlockThreshold, HarmCategory, Modality, SafetySetting, getAI, getGenerativeModel, getVertexAI } from '../src';
20+
import {
21+
testConfigs
22+
} from './constants';
23+
24+
describe('Count Tokens', () => {
25+
26+
testConfigs.forEach(testConfig => {
27+
describe(`${testConfig.toString()}`, () => {
28+
29+
it('CountTokens text', async () => {
30+
const generationConfig: GenerationConfig = {
31+
temperature: 0,
32+
topP: 0,
33+
responseMimeType: 'text/plain'
34+
};
35+
36+
const safetySettings: SafetySetting[] = [
37+
{
38+
category: HarmCategory.HARM_CATEGORY_HARASSMENT,
39+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
40+
method: HarmBlockMethod.PROBABILITY
41+
},
42+
{
43+
category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
44+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
45+
method: HarmBlockMethod.SEVERITY
46+
},
47+
{
48+
category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
49+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
50+
},
51+
{
52+
category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
53+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
54+
}
55+
];
56+
57+
const systemInstruction: Content = {
58+
role: 'system',
59+
parts: [
60+
{
61+
text: 'You are a friendly and helpful assistant.'
62+
}
63+
]
64+
};
65+
const model = getGenerativeModel(testConfig.ai, {
66+
model: testConfig.model,
67+
generationConfig,
68+
systemInstruction,
69+
safetySettings
70+
});
71+
72+
const response = await model.countTokens('Why is the sky blue?');
73+
74+
expect(response.totalTokens).to.equal(6);
75+
expect(response.totalBillableCharacters).to.equal(16);
76+
expect(response.promptTokensDetails).to.not.be.null;
77+
expect(response.promptTokensDetails!.length).to.equal(1);
78+
expect(response.promptTokensDetails![0].modality).to.equal(Modality.TEXT);
79+
expect(response.promptTokensDetails![0].tokenCount).to.equal(6);
80+
});
81+
// TODO (dlarocque): Test countTokens() with the following:
82+
// - inline data
83+
// - public storage reference
84+
// - private storage reference (testing auth integration)
85+
// - count tokens
86+
// - JSON schema
87+
});
88+
})
89+
});

‎packages/ai/integration/firebase-config.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
import * as config from '../../../config/project.json';
2+
3+
export const FIREBASE_CONFIG = config;

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

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
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 } from 'chai';
19+
import { Content, GenerationConfig, HarmBlockMethod, HarmBlockThreshold, HarmCategory, Modality, SafetySetting, getGenerativeModel } from '../src';
20+
import {
21+
testConfigs
22+
} from './constants';
23+
24+
// Token counts are only expected to differ by at most this number of tokens.
25+
// Set to 1 for whitespace that is not always present.
26+
const TOKEN_COUNT_DELTA = 1;
27+
28+
describe('Generate Content', () => {
29+
testConfigs.forEach(testConfig => {
30+
describe(`${testConfig.toString()}`, () => {
31+
32+
it('generateContent', async () => {
33+
const generationConfig: GenerationConfig = {
34+
temperature: 0,
35+
topP: 0,
36+
responseMimeType: 'text/plain'
37+
};
38+
39+
const safetySettings: SafetySetting[] = [
40+
{
41+
category: HarmCategory.HARM_CATEGORY_HARASSMENT,
42+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
43+
method: HarmBlockMethod.PROBABILITY
44+
},
45+
{
46+
category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
47+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
48+
method: HarmBlockMethod.SEVERITY
49+
},
50+
{
51+
category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
52+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
53+
},
54+
{
55+
category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
56+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
57+
}
58+
];
59+
60+
const systemInstruction: Content = {
61+
role: 'system',
62+
parts: [
63+
{
64+
text: 'You are a friendly and helpful assistant.'
65+
}
66+
]
67+
};
68+
69+
const model = getGenerativeModel(testConfig.ai, {
70+
model: testConfig.model,
71+
generationConfig,
72+
safetySettings,
73+
systemInstruction
74+
});
75+
76+
const result = await model.generateContent(
77+
'Where is Google headquarters located? Answer with the city name only.'
78+
);
79+
const response = result.response;
80+
81+
const trimmedText = response.text().trim();
82+
expect(trimmedText).to.equal('Mountain View');
83+
84+
expect(response.usageMetadata).to.not.be.null;
85+
expect(response.usageMetadata!.promptTokenCount).to.be.closeTo(
86+
21,
87+
TOKEN_COUNT_DELTA
88+
);
89+
expect(response.usageMetadata!.candidatesTokenCount).to.be.closeTo(
90+
4,
91+
TOKEN_COUNT_DELTA
92+
);
93+
expect(response.usageMetadata!.totalTokenCount).to.be.closeTo(
94+
25,
95+
TOKEN_COUNT_DELTA * 2
96+
);
97+
expect(response.usageMetadata!.promptTokensDetails).to.not.be.null;
98+
expect(response.usageMetadata!.promptTokensDetails!.length).to.equal(1);
99+
expect(response.usageMetadata!.promptTokensDetails![0].modality).to.equal(
100+
Modality.TEXT
101+
);
102+
expect(response.usageMetadata!.promptTokensDetails![0].tokenCount).to.equal(
103+
21
104+
);
105+
expect(response.usageMetadata!.candidatesTokensDetails).to.not.be.null;
106+
expect(response.usageMetadata!.candidatesTokensDetails!.length).to.equal(1);
107+
expect(
108+
response.usageMetadata!.candidatesTokensDetails![0].modality
109+
).to.equal(Modality.TEXT);
110+
expect(
111+
response.usageMetadata!.candidatesTokensDetails![0].tokenCount
112+
).to.be.closeTo(4, TOKEN_COUNT_DELTA);
113+
});
114+
// TODO (dlarocque): Test generateContentStream
115+
})
116+
})
117+
});

‎packages/vertexai/integration/constants.ts

Lines changed: 0 additions & 63 deletions
This file was deleted.

0 commit comments

Comments
 (0)