Genkit MCP

/** * Copyright 2024 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 { ModelInfo } from 'genkit/model'; import { describe, it } from 'node:test'; import { __setFakeDerivedParams } from '../src/common/index.js'; import { GENERIC_GEMINI_MODEL, gemini } from '../src/gemini.js'; import vertexAI, { gemini15Flash, gemini15Pro } from '../src/index.js'; describe('plugin', () => { __setFakeDerivedParams({ projectId: 'test', location: 'us-central1', }); it('should init the plugin without requiring the api key', async () => { const ai = genkit({ plugins: [vertexAI()], }); assert.ok(ai); }); it('should pre-register a few flagship models', async () => { const ai = genkit({ plugins: [vertexAI()], }); assert.ok(await ai.registry.lookupAction(`/model/${gemini15Flash.name}`)); assert.ok(await ai.registry.lookupAction(`/model/${gemini15Pro.name}`)); }); it('allow referencing models using `gemini` helper', async () => { const ai = genkit({ plugins: [vertexAI()], }); const pro = await ai.registry.lookupAction( `/model/${gemini('gemini-1.5-pro').name}` ); assert.ok(pro); assert.strictEqual(pro.__action.name, 'vertexai/gemini-1.5-pro'); const flash = await ai.registry.lookupAction( `/model/${gemini('gemini-1.5-flash').name}` ); assert.ok(flash); assert.strictEqual(flash.__action.name, 'vertexai/gemini-1.5-flash'); }); it('references explicitly registered models', async () => { const flash002Ref = gemini('gemini-1.5-flash-002'); const ai = genkit({ plugins: [ vertexAI({ location: 'us-central1', models: ['gemini-1.5-pro-002', flash002Ref, 'gemini-4.0-banana'], }), ], }); const pro002Ref = gemini('gemini-1.5-pro-002'); assert.strictEqual(pro002Ref.name, 'vertexai/gemini-1.5-pro-002'); assertEqualModelInfo( pro002Ref.info!, 'Vertex AI - gemini-1.5-pro-002', gemini15Pro.info! ); const pro002 = await ai.registry.lookupAction(`/model/${pro002Ref.name}`); assert.ok(pro002); assert.strictEqual(pro002.__action.name, 'vertexai/gemini-1.5-pro-002'); assertEqualModelInfo( pro002.__action.metadata?.model, 'Vertex AI - gemini-1.5-pro-002', gemini15Pro.info! ); assert.strictEqual(flash002Ref.name, 'vertexai/gemini-1.5-flash-002'); assertEqualModelInfo( flash002Ref.info!, 'Vertex AI - gemini-1.5-flash-002', gemini15Flash.info! ); const flash002 = await ai.registry.lookupAction( `/model/${flash002Ref.name}` ); assert.ok(flash002); assert.strictEqual(flash002.__action.name, 'vertexai/gemini-1.5-flash-002'); assertEqualModelInfo( flash002.__action.metadata?.model, 'Vertex AI - gemini-1.5-flash-002', gemini15Flash.info! ); const bananaRef = gemini('gemini-4.0-banana'); assert.strictEqual(bananaRef.name, 'vertexai/gemini-4.0-banana'); assertEqualModelInfo( bananaRef.info!, 'Vertex AI - gemini-4.0-banana', GENERIC_GEMINI_MODEL.info! // <---- generic model fallback ); const banana = await ai.registry.lookupAction(`/model/${bananaRef.name}`); assert.ok(banana); assert.strictEqual(banana.__action.name, 'vertexai/gemini-4.0-banana'); assertEqualModelInfo( banana.__action.metadata?.model, 'Vertex AI - gemini-4.0-banana', GENERIC_GEMINI_MODEL.info! // <---- generic model fallback ); // this one is not registered const flash003Ref = gemini('gemini-1.5-flash-003'); assert.strictEqual(flash003Ref.name, 'vertexai/gemini-1.5-flash-003'); const flash003 = await ai.registry.lookupAction( `/model/${flash003Ref.name}` ); assert.ok(flash003 === undefined); }); }); function assertEqualModelInfo( modelAction: ModelInfo, expectedLabel: string, expectedInfo: ModelInfo ) { assert.strictEqual(modelAction.label, expectedLabel); assert.deepStrictEqual(modelAction.supports, expectedInfo.supports); assert.deepStrictEqual(modelAction.versions, expectedInfo.versions); }