Skip to content

test(ai): add integration tests #8853

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 30 commits into from
Jun 5, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
98574c2
Skeleton
dlarocque Mar 18, 2025
ae05e53
Integration tests
dlarocque Mar 19, 2025
d83b22e
Import firebase config from secret file
dlarocque Mar 19, 2025
f78282a
Validate that config file exists.
dlarocque Mar 19, 2025
0d1fb41
Cleanup
dlarocque Mar 19, 2025
b083d00
Format
dlarocque Mar 19, 2025
4d6a5ce
Lint fix
dlarocque Mar 19, 2025
3fbb6d6
Extend preprocessor in vertex ai karma conf
dlarocque Mar 19, 2025
d446c77
format
dlarocque Mar 19, 2025
fe6528b
Add placeholder `integration/firebase-config.ts`
dlarocque Mar 19, 2025
7ad7833
Format
dlarocque Mar 20, 2025
11a8b05
Merge branch 'main' into dl/vertex-integration
dlarocque May 16, 2025
9f908d2
Add test configs to run against each test
dlarocque May 16, 2025
abe830c
Add test configs to run against each test
dlarocque May 16, 2025
472c452
count tokens tests
dlarocque May 16, 2025
681b4bd
Format and fix lint issues
dlarocque May 16, 2025
62dab2b
Merge branch 'main' into dl/vertex-integration
dlarocque May 20, 2025
8f65c98
Merge branch 'main' into dl/vertex-integration
dlarocque May 20, 2025
c769317
count tokens tests against both backends
dlarocque May 21, 2025
d78fc29
count tokens test cleanup
dlarocque May 21, 2025
676e90d
remaining integration tests
dlarocque May 22, 2025
ddfd7f5
revert changes to gitignore
dlarocque May 22, 2025
46e910d
Merge branch 'main' into dl/vertex-integration
dlarocque May 22, 2025
3bfcfa7
Merge branch 'dl/vertex-integration' of https://github.com/firebase/f…
dlarocque May 22, 2025
28f843c
use ci config file
dlarocque May 22, 2025
f851be4
import from minified @firebase/ai in integration tests
dlarocque May 23, 2025
f327b4b
exclude integration from rollup
dlarocque May 23, 2025
604a419
revert change to import from @firebase/ai
dlarocque May 23, 2025
9b97310
cleanup
dlarocque May 23, 2025
63f7215
format
dlarocque May 23, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -103,4 +103,4 @@ vertexai-sdk-test-data
mocks-lookup.ts

# temp changeset output
changeset-temp.json
changeset-temp.json
22 changes: 21 additions & 1 deletion .vscode/launch.json
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
"type": "node",
"request": "launch",
"program": "${workspaceFolder}/node_modules/.bin/_mocha",
"cwd": "${workspaceRoot}/packages/vertexai",
"cwd": "${workspaceRoot}/packages/ai",
"args": [
"--require",
"ts-node/register",
Expand All @@ -24,6 +24,26 @@
},
"sourceMaps": true
},
{
"name": "AI Integration Tests (node)",
"type": "node",
"request": "launch",
"program": "${workspaceFolder}/node_modules/.bin/_mocha",
"cwd": "${workspaceRoot}/packages/ai",
"args": [
"--require",
"ts-node/register",
"--require",
"src/index.node.ts",
"--timeout",
"5000",
"integration/**/*.test.ts"
],
"env": {
"TS_NODE_COMPILER_OPTIONS": "{\"module\":\"commonjs\"}"
},
"sourceMaps": true
},
{
"type": "node",
"request": "launch",
Expand Down
132 changes: 132 additions & 0 deletions packages/ai/integration/chat.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
/**
* @license
* Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

import { expect } from 'chai';
import {
Content,
GenerationConfig,
HarmBlockThreshold,
HarmCategory,
SafetySetting,
getGenerativeModel
} from '../src';
import { testConfigs, TOKEN_COUNT_DELTA } from './constants';

describe('Chat Session', () => {
testConfigs.forEach(testConfig => {
describe(`${testConfig.toString()}`, () => {
const commonGenerationConfig: GenerationConfig = {
temperature: 0,
topP: 0,
responseMimeType: 'text/plain'
};

const commonSafetySettings: SafetySetting[] = [
{
category: HarmCategory.HARM_CATEGORY_HARASSMENT,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
},
{
category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
},
{
category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
},
{
category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
}
];

const commonSystemInstruction: Content = {
role: 'system',
parts: [
{
text: 'You are a friendly and helpful assistant.'
}
]
};

it('startChat and sendMessage: text input, text output', async () => {
const model = getGenerativeModel(testConfig.ai, {
model: testConfig.model,
generationConfig: commonGenerationConfig,
safetySettings: commonSafetySettings,
systemInstruction: commonSystemInstruction
});

const chat = model.startChat();
const result1 = await chat.sendMessage(
'What is the capital of France?'
);
const response1 = result1.response;
expect(response1.text().trim().toLowerCase()).to.include('paris');

let history = await chat.getHistory();
expect(history.length).to.equal(2);
expect(history[0].role).to.equal('user');
expect(history[0].parts[0].text).to.equal(
'What is the capital of France?'
);
expect(history[1].role).to.equal('model');
expect(history[1].parts[0].text?.toLowerCase()).to.include('paris');

expect(response1.usageMetadata).to.not.be.null;
// Token counts can vary slightly in chat context
expect(response1.usageMetadata!.promptTokenCount).to.be.closeTo(
15, // "What is the capital of France?" + system instruction
TOKEN_COUNT_DELTA + 2 // More variance for chat context
);
expect(response1.usageMetadata!.candidatesTokenCount).to.be.closeTo(
8, // "Paris"
TOKEN_COUNT_DELTA
);
expect(response1.usageMetadata!.totalTokenCount).to.be.closeTo(
23, // "What is the capital of France?" + system instruction + "Paris"
TOKEN_COUNT_DELTA + 3 // More variance for chat context
);

const result2 = await chat.sendMessage('And what about Italy?');
const response2 = result2.response;
expect(response2.text().trim().toLowerCase()).to.include('rome');

history = await chat.getHistory();
expect(history.length).to.equal(4);
expect(history[2].role).to.equal('user');
expect(history[2].parts[0].text).to.equal('And what about Italy?');
expect(history[3].role).to.equal('model');
expect(history[3].parts[0].text?.toLowerCase()).to.include('rome');

expect(response2.usageMetadata).to.not.be.null;
expect(response2.usageMetadata!.promptTokenCount).to.be.closeTo(
28, // History + "And what about Italy?" + system instruction
TOKEN_COUNT_DELTA + 5 // More variance for chat context with history
);
expect(response2.usageMetadata!.candidatesTokenCount).to.be.closeTo(
8,
TOKEN_COUNT_DELTA
);
expect(response2.usageMetadata!.totalTokenCount).to.be.closeTo(
36,
TOKEN_COUNT_DELTA
);
});
});
});
});
81 changes: 81 additions & 0 deletions packages/ai/integration/constants.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
/**
* @license
* Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

import { initializeApp } from '@firebase/app';
import {
AI,
Backend,
BackendType,
GoogleAIBackend,
VertexAIBackend,
getAI
} from '../src';
import { FIREBASE_CONFIG } from './firebase-config';

const app = initializeApp(FIREBASE_CONFIG);

/**
* Test config that all tests will be ran against.
*/
export type TestConfig = Readonly<{
ai: AI;
model: string;
/** This will be used to output the test config at runtime */
toString: () => string;
}>;

function formatConfigAsString(config: { ai: AI; model: string }): string {
return `${backendNames.get(config.ai.backend.backendType)} ${config.model}`;
}

const backends: readonly Backend[] = [
new GoogleAIBackend(),
new VertexAIBackend()
];

const backendNames: Map<BackendType, string> = new Map([
[BackendType.GOOGLE_AI, 'Google AI'],
[BackendType.VERTEX_AI, 'Vertex AI']
]);

const modelNames: readonly string[] = ['gemini-2.0-flash'];

/**
* Array of test configurations that is iterated over to get full coverage
* of backends and models. Contains all combinations of backends and models.
*/
export const testConfigs: readonly TestConfig[] = backends.flatMap(backend => {
return modelNames.map(modelName => {
const ai = getAI(app, { backend });
return {
ai: getAI(app, { backend }),
model: modelName,
toString: () => formatConfigAsString({ ai, model: modelName })
};
});
});

export const TINY_IMG_BASE64 =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABAQMAAAAl21bKAAAAA1BMVEUAAACnej3aAAAAAXRSTlMAQObYZgAAAApJREFUCNdjYAAAAAIAAeIhvDMAAAAASUVORK5CYII=';
export const IMAGE_MIME_TYPE = 'image/png';
export const TINY_MP3_BASE64 =
'SUQzBAAAAAAAIlRTU0UAAAAOAAADTGF2ZjYxLjcuMTAwAAAAAAAAAAAAAAD/+0DAAAAAAAAAAAAAAAAAAAAAAABJbmZvAAAADwAAAAUAAAK+AGhoaGhoaGhoaGhoaGhoaGhoaGiOjo6Ojo6Ojo6Ojo6Ojo6Ojo6OjrS0tLS0tLS0tLS0tLS0tLS0tLS02tra2tra2tra2tra2tra2tra2tr//////////////////////////wAAAABMYXZjNjEuMTkAAAAAAAAAAAAAAAAkAwYAAAAAAAACvhC6DYoAAAAAAP/7EMQAA8AAAaQAAAAgAAA0gAAABExBTUUzLjEwMFVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVV//sQxCmDwAABpAAAACAAADSAAAAEVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVX/+xDEUwPAAAGkAAAAIAAANIAAAARVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVf/7EMR8g8AAAaQAAAAgAAA0gAAABFVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVV//sQxKYDwAABpAAAACAAADSAAAAEVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVU=';
export const AUDIO_MIME_TYPE = 'audio/mpeg';

// Token counts are only expected to differ by at most this number of tokens.
// Set to 1 for whitespace that is not always present.
export const TOKEN_COUNT_DELTA = 1;
Loading
Loading