test(ai): Add all general use models to integration tests (#9413)

Author: dlarocque

packages/ai/integration/chat.test.ts

@@ -26,7 +26,8 @@ import {
 } from '../src';
 import { testConfigs } from './constants';
 
-describe('Chat Session', () => {
+describe('Chat Session', function () {
+  this.timeout(20_000);
   testConfigs.forEach(testConfig => {
     describe(`${testConfig.toString()}`, () => {
       const commonGenerationConfig: GenerationConfig = {

packages/ai/integration/constants.ts

@@ -52,7 +52,14 @@ const backendNames: Map<BackendType, string> = new Map([
   [BackendType.VERTEX_AI, 'Vertex AI']
 ]);
 
-const modelNames: readonly string[] = ['gemini-2.0-flash', 'gemini-2.5-flash'];
+const modelNames: readonly string[] = [
+  'gemini-2.0-flash-001',
+  'gemini-2.0-flash-lite-001',
+  'gemini-2.5-flash',
+  'gemini-2.5-flash-lite',
+  'gemini-2.5-pro',
+  'gemini-3-pro-preview'
+];
 
 // The Live API requires a different set of models, and they're different for each backend.
 const liveModelNames: Map<BackendType, string[]> = new Map([

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

@@ -118,9 +118,18 @@ describe('Count Tokens', () => {
         };
         const response = await model.countTokens([imagePart]);
 
+        let expectedImageTokens: number;
+        if (testConfig.model === 'gemini-3-pro-preview') {
+          expectedImageTokens =
+            testConfig.ai.backend.backendType === BackendType.GOOGLE_AI
+              ? 1089
+              : 1120;
+        } else {
+          expectedImageTokens = 258;
+        }
+
         if (testConfig.ai.backend.backendType === BackendType.GOOGLE_AI) {
-          const expectedImageTokens = 259;
-          expect(response.totalTokens).to.equal(expectedImageTokens);
+          expect(response.totalTokens).to.equal(expectedImageTokens + 1); // There will be 1 unexpected text token
           expect(response.totalBillableCharacters).to.be.undefined; // Incorrect behavior
           expect(response.promptTokensDetails!.length).to.equal(2);
           expect(response.promptTokensDetails![0]).to.deep.equal({
@@ -129,19 +138,18 @@ describe('Count Tokens', () => {
           });
           expect(response.promptTokensDetails![1]).to.deep.equal({
             modality: Modality.IMAGE,
-            tokenCount: 258
+            tokenCount: expectedImageTokens
           });
         } else if (
           testConfig.ai.backend.backendType === BackendType.VERTEX_AI
         ) {
-          const expectedImageTokens = 258;
           expect(response.totalTokens).to.equal(expectedImageTokens);
           expect(response.totalBillableCharacters).to.be.undefined; // Incorrect behavior
           expect(response.promptTokensDetails!.length).to.equal(1);
           // Note: No text tokens are present for Vertex AI with image-only input.
           expect(response.promptTokensDetails![0]).to.deep.equal({
             modality: Modality.IMAGE,
-            tokenCount: 258
+            tokenCount: expectedImageTokens
           });
           expect(response.promptTokensDetails![0].tokenCount).to.equal(
             expectedImageTokens
@@ -220,13 +228,23 @@ describe('Count Tokens', () => {
         expect(response.promptTokensDetails).to.exist;
         expect(response.promptTokensDetails!.length).to.equal(3);
 
+        let expectedImageTokenCount;
+        if (testConfig.model === 'gemini-3-pro-preview') {
+          expectedImageTokenCount =
+            testConfig.ai.backend.backendType === BackendType.GOOGLE_AI
+              ? 1089
+              : 1120;
+        } else {
+          expectedImageTokenCount = 258;
+        }
+
         expect(imageDetails).to.deep.equal({
           modality: Modality.IMAGE,
-          tokenCount: 258
+          tokenCount: expectedImageTokenCount
         });
 
         if (testConfig.ai.backend.backendType === BackendType.GOOGLE_AI) {
-          expect(response.totalTokens).to.equal(267);
+          expect(response.totalTokens).to.equal(expectedImageTokenCount + 9);
           expect(response.totalBillableCharacters).to.be.undefined;
           expect(textDetails).to.deep.equal({
             modality: Modality.TEXT,
@@ -239,7 +257,7 @@ describe('Count Tokens', () => {
         } else if (
           testConfig.ai.backend.backendType === BackendType.VERTEX_AI
         ) {
-          expect(response.totalTokens).to.equal(261);
+          expect(response.totalTokens).to.equal(expectedImageTokenCount + 3);
           expect(textDetails).to.deep.equal({
             modality: Modality.TEXT,
             tokenCount: 3
@@ -269,7 +287,12 @@ describe('Count Tokens', () => {
 
         const response = await model.countTokens([filePart]);
 
-        const expectedFileTokens = 258;
+        let expectedFileTokens: number;
+        if (testConfig.model === 'gemini-3-pro-preview') {
+          expectedFileTokens = 1120;
+        } else {
+          expectedFileTokens = 258;
+        }
         expect(response.totalTokens).to.equal(expectedFileTokens);
         expect(response.totalBillableCharacters).to.be.undefined;
         expect(response.promptTokensDetails).to.exist;

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

@@ -32,7 +32,7 @@ import {
 import { testConfigs } from './constants';
 
 describe('Generate Content', function () {
-  this.timeout(20_000);
+  this.timeout(90_000); // gemini 3 requests take a long time, especially when using google search and url context.
   testConfigs.forEach(testConfig => {
     describe(`${testConfig.toString()}`, () => {
       const commonGenerationConfig: GenerationConfig = {
@@ -175,8 +175,9 @@ describe('Generate Content', function () {
       describe('URL Context', async () => {
         // URL Context is not supported in Google AI for gemini-2.0-flash
         if (
-          testConfig.ai.backend.backendType === BackendType.GOOGLE_AI &&
-          testConfig.model === 'gemini-2.0-flash'
+          ['gemini-2.0-flash-001', 'gemini-2.0-flash-lite-001'].includes(
+            testConfig.model
+          ) // Models that don't support URL Context
         ) {
           return;
         }
@@ -232,9 +233,7 @@ describe('Generate Content', function () {
           const urlContextMetadata =
             response.candidates?.[0].urlContextMetadata;
           const groundingMetadata = response.candidates?.[0].groundingMetadata;
-          expect(trimmedText).to.contain(
-            'hypermedia information retrieval initiative'
-          );
+          expect(trimmedText.length).to.be.greaterThan(0);
           expect(urlContextMetadata?.urlMetadata).to.exist;
           expect(
             urlContextMetadata?.urlMetadata.length
@@ -302,6 +301,10 @@ describe('Generate Content', function () {
       });
 
       it('generateContent: code execution', async () => {
+        if (testConfig.model === 'gemini-2.0-flash-lite-001') {
+          // This model does not support code execution
+          return;
+        }
         const model = getGenerativeModel(testConfig.ai, {
           model: testConfig.model,
           generationConfig: commonGenerationConfig,