Skip to content

Commit b083d00

Browse files
dlarocquedlarocque
committed
Format
1 parent 0d1fb41 commit b083d00

File tree

4 files changed

+134
-53
lines changed

4 files changed

+134
-53
lines changed

‎packages/vertexai/integration/constants.ts

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,36 @@
1-
import { Content, GenerationConfig, HarmBlockMethod, HarmBlockThreshold, HarmCategory, SafetySetting } from "../src";
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+
*/
217

3-
export const MODEL_NAME = 'gemini-1.5-pro';
18+
import {
19+
Content,
20+
GenerationConfig,
21+
HarmBlockMethod,
22+
HarmBlockThreshold,
23+
HarmCategory,
24+
SafetySetting
25+
} from '../src';
26+
27+
export const MODEL_NAME = 'gemini-1.5-pro';
428

529
export const generationConfig: GenerationConfig = {
630
temperature: 0,
731
topP: 0,
832
responseMimeType: 'text/plain'
9-
}
33+
};
1034

1135
export const safetySettings: SafetySetting[] = [
1236
{
@@ -21,12 +45,12 @@ export const safetySettings: SafetySetting[] = [
2145
},
2246
{
2347
category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
24-
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
48+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
2549
},
2650
{
2751
category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
28-
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
29-
},
52+
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
53+
}
3054
];
3155

3256
export const systemInstruction: Content = {
@@ -36,4 +60,4 @@ export const systemInstruction: Content = {
3660
text: 'You are a friendly and helpful assistant.'
3761
}
3862
]
39-
};
63+
};

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

Lines changed: 36 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,42 @@
1-
import { expect } from "chai";
2-
import { Modality, getGenerativeModel, getVertexAI } from "../src";
3-
import { MODEL_NAME, generationConfig, systemInstruction, safetySettings, } from "./constants";
4-
import { initializeApp } from "@firebase/app";
5-
import { FIREBASE_CONFIG } from "./firebase-config";
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+
*/
617

7-
describe('Count Tokens', () => {
18+
import { expect } from 'chai';
19+
import { Modality, getGenerativeModel, getVertexAI } from '../src';
20+
import {
21+
MODEL_NAME,
22+
generationConfig,
23+
systemInstruction,
24+
safetySettings
25+
} from './constants';
26+
import { initializeApp } from '@firebase/app';
27+
import { FIREBASE_CONFIG } from './firebase-config';
828

9-
before(() => initializeApp(FIREBASE_CONFIG))
29+
describe('Count Tokens', () => {
30+
before(() => initializeApp(FIREBASE_CONFIG));
1031

1132
it('CountTokens text', async () => {
12-
const vertexAI = getVertexAI();
13-
const model = getGenerativeModel(
14-
vertexAI,
15-
{
16-
model: MODEL_NAME,
17-
generationConfig,
18-
systemInstruction,
19-
safetySettings
20-
}
21-
);
33+
const vertexAI = getVertexAI();
34+
const model = getGenerativeModel(vertexAI, {
35+
model: MODEL_NAME,
36+
generationConfig,
37+
systemInstruction,
38+
safetySettings
39+
});
2240

2341
let response = await model.countTokens('Why is the sky blue?');
2442

@@ -35,4 +53,4 @@ describe('Count Tokens', () => {
3553
// - private storage reference (testing auth integration)
3654
// - count tokens
3755
// - JSON schema
38-
});
56+
});

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

Lines changed: 64 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,84 @@
1-
import { expect } from "chai";
2-
import { Modality, getGenerativeModel, getVertexAI } from "../src";
3-
import { MODEL_NAME, generationConfig, systemInstruction, safetySettings } from "./constants";
4-
import { initializeApp } from "@firebase/app";
5-
import { FIREBASE_CONFIG } from "./firebase-config";
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 { Modality, getGenerativeModel, getVertexAI } from '../src';
20+
import {
21+
MODEL_NAME,
22+
generationConfig,
23+
systemInstruction,
24+
safetySettings
25+
} from './constants';
26+
import { initializeApp } from '@firebase/app';
27+
import { FIREBASE_CONFIG } from './firebase-config';
628

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

1133
describe('Generate Content', () => {
12-
13-
before(() => initializeApp(FIREBASE_CONFIG))
34+
before(() => initializeApp(FIREBASE_CONFIG));
1435

1536
it('generateContent', async () => {
16-
const vertexAI = getVertexAI();
17-
const model = getGenerativeModel(
18-
vertexAI,
19-
{
20-
model: MODEL_NAME,
21-
generationConfig,
22-
systemInstruction,
23-
safetySettings
24-
}
25-
);
37+
const vertexAI = getVertexAI();
38+
const model = getGenerativeModel(vertexAI, {
39+
model: MODEL_NAME,
40+
generationConfig,
41+
systemInstruction,
42+
safetySettings
43+
});
2644

27-
const result = await model.generateContent("Where is Google headquarters located? Answer with the city name only.");
45+
const result = await model.generateContent(
46+
'Where is Google headquarters located? Answer with the city name only.'
47+
);
2848
const response = result.response;
29-
49+
3050
const trimmedText = response.text().trim();
3151
expect(trimmedText).to.equal('Mountain View');
3252

3353
expect(response.usageMetadata).to.not.be.null;
34-
expect(response.usageMetadata!.promptTokenCount).to.be.closeTo(21, TOKEN_COUNT_DELTA);
35-
expect(response.usageMetadata!.candidatesTokenCount).to.be.closeTo(4, TOKEN_COUNT_DELTA);
36-
expect(response.usageMetadata!.totalTokenCount).to.be.closeTo(25, TOKEN_COUNT_DELTA*2);
54+
expect(response.usageMetadata!.promptTokenCount).to.be.closeTo(
55+
21,
56+
TOKEN_COUNT_DELTA
57+
);
58+
expect(response.usageMetadata!.candidatesTokenCount).to.be.closeTo(
59+
4,
60+
TOKEN_COUNT_DELTA
61+
);
62+
expect(response.usageMetadata!.totalTokenCount).to.be.closeTo(
63+
25,
64+
TOKEN_COUNT_DELTA * 2
65+
);
3766
expect(response.usageMetadata!.promptTokensDetails).to.not.be.null;
3867
expect(response.usageMetadata!.promptTokensDetails!.length).to.equal(1);
39-
expect(response.usageMetadata!.promptTokensDetails![0].modality).to.equal(Modality.TEXT);
40-
expect(response.usageMetadata!.promptTokensDetails![0].tokenCount).to.equal(21);
68+
expect(response.usageMetadata!.promptTokensDetails![0].modality).to.equal(
69+
Modality.TEXT
70+
);
71+
expect(response.usageMetadata!.promptTokensDetails![0].tokenCount).to.equal(
72+
21
73+
);
4174
expect(response.usageMetadata!.candidatesTokensDetails).to.not.be.null;
4275
expect(response.usageMetadata!.candidatesTokensDetails!.length).to.equal(1);
43-
expect(response.usageMetadata!.candidatesTokensDetails![0].modality).to.equal(Modality.TEXT);
44-
expect(response.usageMetadata!.candidatesTokensDetails![0].tokenCount).to.equal(4);
76+
expect(
77+
response.usageMetadata!.candidatesTokensDetails![0].modality
78+
).to.equal(Modality.TEXT);
79+
expect(
80+
response.usageMetadata!.candidatesTokensDetails![0].tokenCount
81+
).to.equal(4);
4582
});
4683
// TODO (dlarocque): Test generateContentStream
47-
});
84+
});

‎packages/vertexai/karma.conf.js

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@ const files = [`src/**/*.test.ts`];
2424
// Validate that the file that defines the Firebase config to be used in the integration tests exists.
2525
if (argv.integration) {
2626
if (!existsSync('integration/firebase-config.ts')) {
27-
throw new Error(`integration/firebase-config.ts does not exist. This file must contain a Firebase config for a project with Vertex AI enabled.`)
27+
throw new Error(
28+
`integration/firebase-config.ts does not exist. This file must contain a Firebase config for a project with Vertex AI enabled.`
29+
);
2830
}
2931
}
3032

0 commit comments

Comments
 (0)