index_test.ts•19.6 kB
/**
* 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 * as assert from 'assert';
import { genkit } from 'genkit';
import { afterEach, beforeEach, describe, it, mock } from 'node:test';
import {
TEST_ONLY as EMBEDDER_TEST_ONLY,
EmbeddingConfigSchema,
} from '../../src/googleai/embedder.js';
import {
TEST_ONLY as GEMINI_TEST_ONLY,
GeminiConfigSchema,
GeminiTtsConfigSchema,
GemmaConfigSchema,
} from '../../src/googleai/gemini.js';
import {
TEST_ONLY as IMAGEN_TEST_ONLY,
ImagenConfigSchema,
} from '../../src/googleai/imagen.js';
import { googleAI } from '../../src/googleai/index.js';
import { Model } from '../../src/googleai/types.js';
import { MISSING_API_KEY_ERROR } from '../../src/googleai/utils.js';
import {
TEST_ONLY as VEO_TEST_ONLY,
VeoConfigSchema,
} from '../../src/googleai/veo.js';
describe('GoogleAI Plugin', () => {
let originalEnv: NodeJS.ProcessEnv;
let fetchMock: any;
beforeEach(() => {
originalEnv = { ...process.env };
// Set a default API key for most tests
process.env.GOOGLE_API_KEY = 'test-api-key';
// Mock global.fetch for listActions tests
fetchMock = mock.method(global, 'fetch');
});
afterEach(() => {
// Restore environment variables
process.env = originalEnv;
// Restore mocks
mock.restoreAll();
});
describe('Initializer', () => {
it('should pre-register flagship Gemini models', async () => {
const ai = genkit({ plugins: [googleAI()] });
const model1Name = Object.keys(GEMINI_TEST_ONLY.KNOWN_MODELS)[0];
const model1Path = `/model/googleai/${model1Name}`;
const expectedBaseName = `googleai/${model1Name}`;
const model1 = await ai.registry.lookupAction(model1Path);
assert.ok(model1, `${model1Name} should be registered at ${model1Path}`);
assert.strictEqual(model1?.__action.name, expectedBaseName);
});
it('should register all known Gemini models', async () => {
const ai = genkit({ plugins: [googleAI()] });
for (const modelName in GEMINI_TEST_ONLY.KNOWN_MODELS) {
const modelPath = `/model/googleai/${modelName}`;
const expectedBaseName = `googleai/${modelName}`;
const model = await ai.registry.lookupAction(modelPath);
assert.ok(model, `${modelName} should be registered at ${modelPath}`);
assert.strictEqual(model?.__action.name, expectedBaseName);
}
});
it('should register all known Imagen models', async () => {
const ai = genkit({ plugins: [googleAI()] });
for (const modelName in IMAGEN_TEST_ONLY.KNOWN_MODELS) {
const modelPath = `/model/googleai/${modelName}`;
const expectedBaseName = `googleai/${modelName}`;
const model = await ai.registry.lookupAction(modelPath);
assert.ok(model, `${modelName} should be registered at ${modelPath}`);
assert.strictEqual(model?.__action.name, expectedBaseName);
}
});
it('should register all known Veo models', async () => {
const ai = genkit({ plugins: [googleAI()] });
for (const modelName in VEO_TEST_ONLY.KNOWN_MODELS) {
const modelPath = `/background-model/googleai/${modelName}`;
const expectedBaseName = `googleai/${modelName}`;
const model = await ai.registry.lookupAction(modelPath);
assert.ok(model, `${modelName} should be registered at ${modelPath}`);
assert.strictEqual(model?.__action.name, expectedBaseName);
}
});
it('should pre-register flagship Embedder models', async () => {
const ai = genkit({ plugins: [googleAI()] });
const modelKeys = Object.keys(EMBEDDER_TEST_ONLY.KNOWN_MODELS);
if (modelKeys.length > 0) {
const model1Name = modelKeys[0];
const model1Path = `/embedder/googleai/${model1Name}`;
const expectedBaseName = `googleai/${model1Name}`;
const model1 = await ai.registry.lookupAction(model1Path);
assert.ok(
model1,
`${model1Name} should be registered at ${model1Path}`
);
assert.strictEqual(model1?.__action.name, expectedBaseName);
}
});
it('should register all known Embedder models', async () => {
const ai = genkit({ plugins: [googleAI()] });
for (const modelName in EMBEDDER_TEST_ONLY.KNOWN_MODELS) {
const modelPath = `/embedder/googleai/${modelName}`;
const expectedBaseName = `googleai/${modelName}`;
const model = await ai.registry.lookupAction(modelPath);
assert.ok(model, `${modelName} should be registered at ${modelPath}`);
assert.strictEqual(model?.__action.name, expectedBaseName);
}
});
it('should throw from initializer if no API key is provided', async () => {
delete process.env.GOOGLE_API_KEY;
delete process.env.GEMINI_API_KEY;
delete process.env.GOOGLE_GENAI_API_KEY;
const plugin = googleAI({});
await assert.rejects(async () => {
await plugin.init!();
}, MISSING_API_KEY_ERROR);
});
it('should NOT throw from initializer if apiKey is false', async () => {
delete process.env.GOOGLE_API_KEY;
delete process.env.GEMINI_API_KEY;
delete process.env.GOOGLE_GENAI_API_KEY;
const plugin = googleAI({ apiKey: false });
await assert.doesNotReject(async () => {
await plugin.init!();
});
});
});
describe('Resolver via lookupAction', () => {
const testGeminiModelName = 'gemini-custom-pro';
const testGeminiModelPath = `/model/googleai/${testGeminiModelName}`;
const testImagenModelName = 'imagen-custom';
const testImagenModelPath = `/model/googleai/${testImagenModelName}`;
const testVeoModelName = 'veo-custom';
const testVeoModelPath = `/background-model/googleai/${testVeoModelName}`;
const testEmbedderName = 'embedding-custom-001';
const testEmbedderPath = `/embedder/googleai/${testEmbedderName}`;
it('should register a new Gemini model when looked up', async () => {
const ai = genkit({ plugins: [googleAI()] });
const model = await ai.registry.lookupAction(testGeminiModelPath);
assert.ok(
model,
`${testGeminiModelName} should be resolvable and registered`
);
assert.strictEqual(
model?.__action.name,
`googleai/${testGeminiModelName}`
);
});
it('should register a new Imagen model when looked up', async () => {
const ai = genkit({ plugins: [googleAI()] });
const model = await ai.registry.lookupAction(testImagenModelPath);
assert.ok(
model,
`${testImagenModelName} should be resolvable and registered`
);
assert.strictEqual(
model?.__action.name,
`googleai/${testImagenModelName}`
);
});
it('should register a new Veo model when looked up', async () => {
const ai = genkit({ plugins: [googleAI()] });
const model = await ai.registry.lookupAction(testVeoModelPath);
assert.ok(
model,
`${testVeoModelName} should be resolvable and registered`
);
assert.strictEqual(model?.__action.name, `googleai/${testVeoModelName}`);
});
it('should register a new Embedder when looked up', async () => {
const ai = genkit({ plugins: [googleAI()] });
const embedder = await ai.registry.lookupAction(testEmbedderPath);
assert.ok(
embedder,
`${testEmbedderName} should be resolvable and registered`
);
assert.strictEqual(
embedder?.__action.name,
`googleai/${testEmbedderName}`
);
});
});
describe('googleAI.model', () => {
it('should return a gemini ModelReference with correct schema', () => {
const modelName = 'gemini-2.0-flash';
const modelRef = googleAI.model(modelName);
assert.strictEqual(
modelRef.name,
`googleai/${modelName}`,
'Name should be prefixed'
);
assert.ok(
modelRef.info?.supports?.multiturn,
'Gemini model should support multiturn'
);
assert.strictEqual(
modelRef.configSchema,
GeminiConfigSchema,
'Should have GeminiConfigSchema'
);
});
it('should return a TTS model reference with correct schema', () => {
const modelRef = googleAI.model('gemini-5.0-tts');
assert.strictEqual(
modelRef.configSchema,
GeminiTtsConfigSchema,
'Should have GeminiTTsConfigSchema'
);
assert.ok(
!modelRef.info?.supports?.multiturn,
'Gemini TTS model should not support multiturn'
);
});
it('should have config values for gemini TTS', () => {
const modelRef = googleAI.model('gemini-5.0-tts', {
speechConfig: {
voiceConfig: {
prebuiltVoiceConfig: { voiceName: 'Algenib' },
},
},
});
assert.strictEqual(
modelRef.configSchema,
GeminiTtsConfigSchema,
'Should have GeminiTTsConfigSchema'
);
assert.strictEqual(
modelRef.config?.speechConfig?.voiceConfig?.prebuiltVoiceConfig
?.voiceName,
'Algenib'
);
});
it('should return a Gemma model reference with correct schema', () => {
const modelRef = googleAI.model('gemma-new-model');
assert.strictEqual(
modelRef.configSchema,
GemmaConfigSchema,
'Should have GemmaConfigSchema'
);
assert.ok(
modelRef.info?.supports?.multiturn,
'Gemma model should support multiturn'
);
});
it('should have config values for gemma', () => {
const modelRef = googleAI.model('gemma-3-12b-it', {
temperature: 0.7,
});
assert.strictEqual(modelRef.name, 'googleai/gemma-3-12b-it');
assert.strictEqual(modelRef.config?.temperature, 0.7);
});
it('should return an Imagen model reference with correct schema', () => {
const modelRef = googleAI.model('imagen-new-model');
assert.strictEqual(
modelRef.configSchema,
ImagenConfigSchema,
'Should have ImagenConfigSchema'
);
});
it('should have config values for imagen model', () => {
const modelRef = googleAI.model('imagen-new-model', {
numberOfImages: 4,
aspectRatio: '16:9',
});
assert.strictEqual(
modelRef.configSchema,
ImagenConfigSchema,
'Should have ImagenConfigSchema'
);
assert.strictEqual(
modelRef.config?.numberOfImages,
4,
'should have 4 images'
);
assert.strictEqual(
modelRef.config?.aspectRatio,
'16:9',
'should be 16:9'
);
});
it('should return a Veo model reference with correct schema', () => {
const modelRef = googleAI.model('veo-new-model');
assert.strictEqual(
modelRef.configSchema,
VeoConfigSchema,
'Should have VeoConfigSchema'
);
});
it('should have config values for veo model', () => {
const modelRef = googleAI.model('veo-new-model', {
aspectRatio: '9:16',
durationSeconds: 8,
});
assert.strictEqual(
modelRef.configSchema,
VeoConfigSchema,
'Should have VeoConfigSchema'
);
assert.strictEqual(
modelRef.config?.aspectRatio,
'9:16',
'should be 9:16'
);
assert.strictEqual(
modelRef.config?.durationSeconds,
8,
'should be 8 seconds'
);
});
it('should return a gemini model reference for unknown model names', () => {
const modelRef = googleAI.model('foo-model');
assert.strictEqual(
modelRef.configSchema,
GeminiConfigSchema,
'Should have GeminiConfigSchema'
);
});
it('should have gemini config values for unknown model', () => {
const modelRef = googleAI.model('foo-model', { temperature: 0.3 });
assert.strictEqual(
modelRef.configSchema,
GeminiConfigSchema,
'Should have GeminiConfigSchema'
);
assert.strictEqual(modelRef.config?.temperature, 0.3);
});
it('should handle names with googleai/ prefix', () => {
const modelName = 'googleai/gemini-2.0-pro';
const modelRef = googleAI.model(modelName);
assert.strictEqual(modelRef.name, modelName);
});
it('should handle names with models/ prefix', () => {
const modelName = 'models/gemini-2.0-pro';
const modelRef = googleAI.model(modelName);
assert.strictEqual(modelRef.name, 'googleai/gemini-2.0-pro');
});
});
describe('googleAI.embedder', () => {
it('should return an EmbedderReference with correct schema', () => {
const embedderName = 'text-embedding-004';
const embedderRef = googleAI.embedder(embedderName);
assert.strictEqual(embedderRef.name, `googleai/${embedderName}`);
assert.ok(embedderRef.info, 'Should have info');
assert.strictEqual(
embedderRef.configSchema,
EmbeddingConfigSchema,
'Should have GoogleAIEmbeddingConfigSchema'
);
});
it('should handle names with googleai/ prefix', () => {
const embedderName = 'googleai/text-embedding-custom';
const embedderRef = googleAI.embedder(embedderName);
assert.strictEqual(embedderRef.name, embedderName);
});
it('should handle names with embedders/ prefix', () => {
const embedderName = 'embedders/text-embedding-custom';
const embedderRef = googleAI.embedder(embedderName);
assert.strictEqual(embedderRef.name, 'googleai/text-embedding-custom');
});
});
describe('listActions Function', () => {
const createMockResponse = (models: Array<Partial<Model>>) => {
const responseBody = { models: models };
return Promise.resolve({
ok: true,
status: 200,
json: async () => responseBody,
text: async () => JSON.stringify(responseBody),
});
};
it('should return an empty array if no models are returned', async () => {
fetchMock.mock.mockImplementation(async () => createMockResponse([]));
const plugin = googleAI();
const actions = await plugin.list!();
assert.deepStrictEqual(actions, [], 'Should return an empty array');
});
it('should return metadata for models and embedders', async () => {
const mockModels: Partial<Model>[] = [
{
name: 'models/gemini-2.5-pro',
supportedGenerationMethods: ['generateContent'],
},
{
name: 'models/text-embedding-004',
supportedGenerationMethods: ['embedContent'],
},
{
name: 'models/imagen-3.0-generate-002',
supportedGenerationMethods: ['predict'],
},
{
name: 'models/veo-2.0-generate-001',
supportedGenerationMethods: ['predictLongRunning'],
},
{
name: 'models/other-model',
supportedGenerationMethods: ['other'],
},
];
fetchMock.mock.mockImplementation(async () =>
createMockResponse(mockModels)
);
const plugin = googleAI();
const actions = await plugin.list!();
const actionNames = actions.map((a) => a.name).sort();
assert.deepStrictEqual(
actionNames,
[
'googleai/gemini-2.5-pro',
'googleai/imagen-3.0-generate-002',
'googleai/text-embedding-004',
'googleai/veo-2.0-generate-001',
].sort()
);
const modelAction = actions.find(
(a) => a.name === 'googleai/gemini-2.5-pro'
);
assert.strictEqual(modelAction?.actionType, 'model');
const embedderAction = actions.find(
(a) => a.name === 'googleai/text-embedding-004'
);
assert.strictEqual(embedderAction?.actionType, 'embedder');
const imagenAction = actions.find(
(a) => a.name === 'googleai/imagen-3.0-generate-002'
);
assert.strictEqual(imagenAction?.actionType, 'model');
const veoAction = actions.find(
(a) => a.name === 'googleai/veo-2.0-generate-001'
);
assert.strictEqual(veoAction?.actionType, 'model');
});
it('should filter out deprecated models', async () => {
const mockModels = [
{
name: 'models/gemini-1.5-flash',
supportedGenerationMethods: ['generateContent'],
},
{
name: 'models/gemini-pro-deprecated',
supportedGenerationMethods: ['generateContent'],
description: 'This model is deprecated.',
},
{
name: 'models/imagen-deprecated',
supportedGenerationMethods: ['predict'],
description: 'This model is deprecated.',
},
{
name: 'models/veo-deprecated',
supportedGenerationMethods: ['predictLongRunning'],
description: 'This model is deprecated.',
},
{
name: 'models/text-embedding-deprecated',
supportedGenerationMethods: ['embedContent'],
description: 'deprecated',
},
];
fetchMock.mock.mockImplementation(async () =>
createMockResponse(mockModels)
);
const plugin = googleAI();
const actions = await plugin.list!();
const actionNames = actions.map((a) => a.name);
assert.deepStrictEqual(actionNames, ['googleai/gemini-1.5-flash']);
});
it('should handle fetch errors gracefully', async () => {
fetchMock.mock.mockImplementation(async () => {
return Promise.resolve({
ok: false,
status: 500,
statusText: 'Internal Error',
json: async () => ({ error: { message: 'API Error' } }),
text: async () => JSON.stringify({ error: { message: 'API Error' } }),
});
});
const plugin = googleAI();
const actions = await plugin.list!();
assert.deepStrictEqual(
actions,
[],
'Should return empty array on fetch error'
);
});
it('should return empty array if API key is missing for listActions', async () => {
delete process.env.GOOGLE_API_KEY;
delete process.env.GEMINI_API_KEY;
delete process.env.GOOGLE_GENAI_API_KEY;
const plugin = googleAI({ apiKey: false });
const actions = await plugin.list!();
assert.deepStrictEqual(
actions,
[],
'Should return empty array if API key is not found'
);
assert.strictEqual(
fetchMock.mock.callCount(),
0,
'Fetch should not be called'
);
});
it('should use listActions cache', async () => {
const mockModels = [
{
name: 'models/gemini-1.0-pro',
supportedGenerationMethods: ['generateContent'],
},
];
fetchMock.mock.mockImplementation(async () =>
createMockResponse(mockModels)
);
const plugin = googleAI();
await plugin.list!();
await plugin.list!();
assert.strictEqual(
fetchMock.mock.callCount(),
1,
'fetch should only be called once'
);
});
});
});