diff --git a/js/plugins/google-genai/src/googleai/embedder.ts b/js/plugins/google-genai/src/googleai/embedder.ts index ad5adb4256..100e3a391b 100644 --- a/js/plugins/google-genai/src/googleai/embedder.ts +++ b/js/plugins/google-genai/src/googleai/embedder.ts @@ -20,10 +20,10 @@ import { embedderActionMetadata, EmbedderInfo, EmbedderReference, - Genkit, z, } from 'genkit'; import { embedderRef } from 'genkit/embedder'; +import { embedder as pluginEmbedder } from 'genkit/plugin'; import { embedContent } from './client.js'; import { EmbedContentRequest, @@ -122,43 +122,40 @@ export function listActions(models: Model[]): ActionMetadata[] { ); } -export function defineKnownModels(ai: Genkit, options?: GoogleAIPluginOptions) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineEmbedder(ai, name, options); - } +export function listKnownModels(options?: GoogleAIPluginOptions) { + return Object.keys(KNOWN_MODELS).map((name) => defineEmbedder(name, options)); } export function defineEmbedder( - ai: Genkit, name: string, pluginOptions?: GoogleAIPluginOptions ): EmbedderAction { checkApiKey(pluginOptions?.apiKey); const ref = model(name); - return ai.defineEmbedder( + return pluginEmbedder( { name: ref.name, configSchema: ref.configSchema, info: ref.info, }, - async (input, reqOptions) => { + async (request, _) => { const embedApiKey = calculateApiKey( pluginOptions?.apiKey, - reqOptions?.apiKey + request.options?.apiKey ); - const embedVersion = reqOptions?.version || extractVersion(ref); + const embedVersion = request.options?.version || extractVersion(ref); const embeddings = await Promise.all( - input.map(async (doc) => { + request.input.map(async (doc) => { const response = await embedContent(embedApiKey, embedVersion, { - taskType: reqOptions?.taskType, - title: reqOptions?.title, + taskType: request.options?.taskType, + title: request.options?.title, content: { role: '', parts: [{ text: doc.text }], }, - outputDimensionality: reqOptions?.outputDimensionality, + outputDimensionality: request.options?.outputDimensionality, } as EmbedContentRequest); const values = response.embedding.values; return { embedding: values }; diff --git a/js/plugins/google-genai/src/googleai/gemini.ts b/js/plugins/google-genai/src/googleai/gemini.ts index 7a1118ff60..9d7195b3f0 100644 --- a/js/plugins/google-genai/src/googleai/gemini.ts +++ b/js/plugins/google-genai/src/googleai/gemini.ts @@ -14,13 +14,7 @@ * limitations under the License. */ -import { - ActionMetadata, - Genkit, - GenkitError, - modelActionMetadata, - z, -} from 'genkit'; +import { ActionMetadata, GenkitError, modelActionMetadata, z } from 'genkit'; import { GenerationCommonConfigDescriptions, GenerationCommonConfigSchema, @@ -32,6 +26,7 @@ import { modelRef, } from 'genkit/model'; import { downloadRequestMedia } from 'genkit/model/middleware'; +import { model as pluginModel } from 'genkit/plugin'; import { runInNewSpan } from 'genkit/tracing'; import { fromGeminiCandidate, @@ -434,17 +429,16 @@ export function listActions(models: Model[]): ActionMetadata[] { ); } -export function defineKnownModels(ai: Genkit, options?: GoogleAIPluginOptions) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineModel(ai, name, options); - } +export function listKnownModels(options?: GoogleAIPluginOptions) { + return Object.keys(KNOWN_MODELS).map((name: string) => + defineModel(name, options) + ); } /** * Defines a new GoogleAI Gemini model. */ export function defineModel( - ai: Genkit, name: string, pluginOptions?: GoogleAIPluginOptions ): ModelAction { @@ -482,9 +476,8 @@ export function defineModel( ); } - return ai.defineModel( + return pluginModel( { - apiVersion: 'v2', name: ref.name, ...ref.info, configSchema: ref.configSchema, @@ -660,7 +653,6 @@ export function defineModel( // API params as for input. return pluginOptions?.experimental_debugTraces ? await runInNewSpan( - ai.registry, { metadata: { name: streamingRequested ? 'sendMessageStream' : 'sendMessage', diff --git a/js/plugins/google-genai/src/googleai/imagen.ts b/js/plugins/google-genai/src/googleai/imagen.ts index fdf729a7ba..2d402841ba 100644 --- a/js/plugins/google-genai/src/googleai/imagen.ts +++ b/js/plugins/google-genai/src/googleai/imagen.ts @@ -20,7 +20,6 @@ import { MessageData, modelActionMetadata, z, - type Genkit, } from 'genkit'; import { getBasicUsageStats, @@ -30,6 +29,7 @@ import { type ModelInfo, type ModelReference, } from 'genkit/model'; +import { model as pluginModel } from 'genkit/plugin'; import { imagenPredict } from './client.js'; import type { ClientOptions, @@ -169,14 +169,13 @@ export function listActions(models: Model[]): ActionMetadata[] { }); } -export function defineKnownModels(ai: Genkit, options?: GoogleAIPluginOptions) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineModel(ai, name, options); - } +export function listKnownModels(options?: GoogleAIPluginOptions) { + return Object.keys(KNOWN_MODELS).map((name: string) => + defineModel(name, options) + ); } export function defineModel( - ai: Genkit, name: string, pluginOptions?: GoogleAIPluginOptions ): ModelAction { @@ -187,9 +186,8 @@ export function defineModel( baseUrl: pluginOptions?.baseUrl, }; - return ai.defineModel( + return pluginModel( { - apiVersion: 'v2', name: ref.name, ...ref.info, configSchema: ref.configSchema, diff --git a/js/plugins/google-genai/src/googleai/index.ts b/js/plugins/google-genai/src/googleai/index.ts index bbd5ea9647..220854ce04 100644 --- a/js/plugins/google-genai/src/googleai/index.ts +++ b/js/plugins/google-genai/src/googleai/index.ts @@ -14,15 +14,13 @@ * limitations under the License. */ -import { - ActionMetadata, - EmbedderReference, - Genkit, - ModelReference, - z, -} from 'genkit'; +import { ActionMetadata, EmbedderReference, ModelReference, z } from 'genkit'; import { logger } from 'genkit/logging'; -import { GenkitPlugin, genkitPlugin } from 'genkit/plugin'; +import { + GenkitPluginV2, + ResolvableAction, + genkitPluginV2, +} from 'genkit/plugin'; import { ActionType } from 'genkit/registry'; import { extractErrMsg } from '../common/utils.js'; import { listModels } from './client.js'; @@ -42,41 +40,41 @@ export { type GeminiConfig, type GeminiTtsConfig } from './gemini.js'; export { type ImagenConfig } from './imagen.js'; export { type GoogleAIPluginOptions }; -async function initializer(ai: Genkit, options?: GoogleAIPluginOptions) { - imagen.defineKnownModels(ai, options); - gemini.defineKnownModels(ai, options); - embedder.defineKnownModels(ai, options); - veo.defineKnownModels(ai, options); +async function initializer(options?: GoogleAIPluginOptions) { + return [ + ...imagen.listKnownModels(options), + ...gemini.listKnownModels(options), + ...embedder.listKnownModels(options), + ...veo.listKnownModels(options), + ]; } async function resolver( - ai: Genkit, actionType: ActionType, actionName: string, options: GoogleAIPluginOptions -) { +): Promise { switch (actionType) { case 'model': if (veo.isVeoModelName(actionName)) { - // no-op (not gemini) + return undefined; } else if (imagen.isImagenModelName(actionName)) { - imagen.defineModel(ai, actionName, options); + return await imagen.defineModel(actionName, options); } else { // gemini, tts, gemma, unknown models - gemini.defineModel(ai, actionName, options); + return await gemini.defineModel(actionName, options); } break; case 'background-model': if (veo.isVeoModelName(actionName)) { - veo.defineModel(ai, actionName, options); + return await veo.defineModel(actionName, options); } break; case 'embedder': - embedder.defineEmbedder(ai, actionName, options); + return await embedder.defineEmbedder(actionName, options); break; - default: - // no-op } + return undefined; } async function listActions( @@ -104,23 +102,25 @@ async function listActions( /** * Google Gemini Developer API plugin. */ -export function googleAIPlugin(options?: GoogleAIPluginOptions): GenkitPlugin { +export function googleAIPlugin( + options?: GoogleAIPluginOptions +): GenkitPluginV2 { let listActionsCache; - return genkitPlugin( - 'googleai', - async (ai: Genkit) => await initializer(ai, options), - async (ai: Genkit, actionType: ActionType, actionName: string) => - await resolver(ai, actionType, actionName, options || {}), - async () => { + return genkitPluginV2({ + name: 'googleai', + init: async () => await initializer(options), + resolve: async (actionType: ActionType, actionName: string) => + await resolver(actionType, actionName, options || {}), + list: async () => { if (listActionsCache) return listActionsCache; listActionsCache = await listActions(options); return listActionsCache; - } - ); + }, + }); } export type GoogleAIPlugin = { - (pluginOptions?: GoogleAIPluginOptions): GenkitPlugin; + (pluginOptions?: GoogleAIPluginOptions): GenkitPluginV2; model( name: gemini.KnownGemmaModels | (gemini.GemmaModelName & {}), config: gemini.GemmaConfig diff --git a/js/plugins/google-genai/src/googleai/veo.ts b/js/plugins/google-genai/src/googleai/veo.ts index 4d6cd41b60..cd05113b0e 100644 --- a/js/plugins/google-genai/src/googleai/veo.ts +++ b/js/plugins/google-genai/src/googleai/veo.ts @@ -20,7 +20,6 @@ import { Operation, modelActionMetadata, z, - type Genkit, } from 'genkit'; import { BackgroundModelAction, @@ -29,6 +28,7 @@ import { type ModelInfo, type ModelReference, } from 'genkit/model'; +import { backgroundModel as pluginBackgroundModel } from 'genkit/plugin'; import { veoCheckOperation, veoPredict } from './client.js'; import { ClientOptions, @@ -156,17 +156,16 @@ export function listActions(models: Model[]): ActionMetadata[] { ); } -export function defineKnownModels(ai: Genkit, options?: GoogleAIPluginOptions) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineModel(ai, name, options); - } +export function listKnownModels(options?: GoogleAIPluginOptions) { + return Object.keys(KNOWN_MODELS).map((name: string) => + defineModel(name, options) + ); } /** * Defines a new GoogleAI Veo model. */ export function defineModel( - ai: Genkit, name: string, pluginOptions?: GoogleAIPluginOptions ): BackgroundModelAction { @@ -176,7 +175,7 @@ export function defineModel( baseUrl: pluginOptions?.baseUrl, }; - return ai.defineBackgroundModel({ + return pluginBackgroundModel({ name: ref.name, ...ref.info, configSchema: ref.configSchema, diff --git a/js/plugins/google-genai/src/vertexai/embedder.ts b/js/plugins/google-genai/src/vertexai/embedder.ts index c7a8001837..cbf5025284 100644 --- a/js/plugins/google-genai/src/vertexai/embedder.ts +++ b/js/plugins/google-genai/src/vertexai/embedder.ts @@ -14,13 +14,14 @@ * limitations under the License. */ -import { z, type Document, type Genkit } from 'genkit'; +import { z, type Document } from 'genkit'; import { EmbedderInfo, embedderRef, type EmbedderAction, type EmbedderReference, } from 'genkit/embedder'; +import { embedder as pluginEmbedder } from 'genkit/plugin'; import { embedContent } from './client.js'; import { ClientOptions, @@ -145,36 +146,36 @@ export function model( }); } -export function defineKnownModels( - ai: Genkit, +export function listKnownModels( clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineEmbedder(ai, name, clientOptions, pluginOptions); - } + return Object.keys(KNOWN_MODELS).map((name) => + defineEmbedder(name, clientOptions, pluginOptions) + ); } export function defineEmbedder( - ai: Genkit, name: string, clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ): EmbedderAction { const ref = model(name); - return ai.defineEmbedder( + return pluginEmbedder( { name: ref.name, configSchema: ref.configSchema, info: ref.info!, }, - async (input, options?: EmbeddingConfig) => { + async (request) => { const embedContentRequest: EmbedContentRequest = { - instances: input.map((doc: Document) => - toEmbeddingInstance(ref, doc, options) + instances: request.input.map((doc: Document) => + toEmbeddingInstance(ref, doc, request.options) ), - parameters: { outputDimensionality: options?.outputDimensionality }, + parameters: { + outputDimensionality: request.options?.outputDimensionality, + }, }; const response = await embedContent( diff --git a/js/plugins/google-genai/src/vertexai/gemini.ts b/js/plugins/google-genai/src/vertexai/gemini.ts index 9ad61fe3e6..66467b81ce 100644 --- a/js/plugins/google-genai/src/vertexai/gemini.ts +++ b/js/plugins/google-genai/src/vertexai/gemini.ts @@ -14,13 +14,7 @@ * limitations under the License. */ -import { - ActionMetadata, - Genkit, - GenkitError, - modelActionMetadata, - z, -} from 'genkit'; +import { ActionMetadata, GenkitError, modelActionMetadata, z } from 'genkit'; import { GenerationCommonConfigDescriptions, GenerationCommonConfigSchema, @@ -32,6 +26,7 @@ import { modelRef, } from 'genkit/model'; import { downloadRequestMedia } from 'genkit/model/middleware'; +import { model as pluginModel } from 'genkit/plugin'; import { runInNewSpan } from 'genkit/tracing'; import { fromGeminiCandidate, @@ -403,21 +398,19 @@ export function listActions(models: Model[]): ActionMetadata[] { }); } -export function defineKnownModels( - ai: Genkit, +export function listKnownModels( clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineModel(ai, name, clientOptions, pluginOptions); - } + return Object.keys(KNOWN_MODELS).map((name) => + defineModel(name, clientOptions, pluginOptions) + ); } /** * Define a Vertex AI Gemini model. */ export function defineModel( - ai: Genkit, name: string, clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions @@ -446,9 +439,8 @@ export function defineModel( ); } - return ai.defineModel( + return pluginModel( { - apiVersion: 'v2', name: ref.name, ...ref.info, configSchema: ref.configSchema, @@ -682,7 +674,6 @@ export function defineModel( const msg = toGeminiMessage(messages[messages.length - 1], ref); return pluginOptions?.experimental_debugTraces ? await runInNewSpan( - ai.registry, { metadata: { name: streamingRequested ? 'sendMessageStream' : 'sendMessage', diff --git a/js/plugins/google-genai/src/vertexai/imagen.ts b/js/plugins/google-genai/src/vertexai/imagen.ts index ab3c0de4f3..6e1c399096 100644 --- a/js/plugins/google-genai/src/vertexai/imagen.ts +++ b/js/plugins/google-genai/src/vertexai/imagen.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { ActionMetadata, Genkit, modelActionMetadata, z } from 'genkit'; +import { ActionMetadata, modelActionMetadata, z } from 'genkit'; import { GenerationCommonConfigSchema, ModelAction, @@ -22,6 +22,7 @@ import { ModelReference, modelRef, } from 'genkit/model'; +import { model as pluginModel } from 'genkit/plugin'; import { imagenPredict } from './client.js'; import { fromImagenResponse, toImagenPredictRequest } from './converters.js'; import { ClientOptions, Model, VertexPluginOptions } from './types.js'; @@ -246,27 +247,24 @@ export function listActions(models: Model[]): ActionMetadata[] { }); } -export function defineKnownModels( - ai: Genkit, +export function listKnownModels( clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineModel(ai, name, clientOptions, pluginOptions); - } + return Object.keys(KNOWN_MODELS).map((name: string) => + defineModel(name, clientOptions, pluginOptions) + ); } export function defineModel( - ai: Genkit, name: string, clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ): ModelAction { const ref = model(name); - return ai.defineModel( + return pluginModel( { - apiVersion: 'v2', name: ref.name, ...ref.info, configSchema: ref.configSchema, diff --git a/js/plugins/google-genai/src/vertexai/index.ts b/js/plugins/google-genai/src/vertexai/index.ts index 6e07f9b593..6b061f12ad 100644 --- a/js/plugins/google-genai/src/vertexai/index.ts +++ b/js/plugins/google-genai/src/vertexai/index.ts @@ -20,8 +20,12 @@ * @module / */ -import { EmbedderReference, Genkit, ModelReference, z } from 'genkit'; -import { GenkitPlugin, genkitPlugin } from 'genkit/plugin'; +import { EmbedderReference, ModelReference, z } from 'genkit'; +import { + GenkitPluginV2, + ResolvableAction, + genkitPluginV2, +} from 'genkit/plugin'; import { ActionType } from 'genkit/registry'; import { listModels } from './client.js'; @@ -41,45 +45,45 @@ export { type LyriaConfig } from './lyria.js'; export { type VertexPluginOptions } from './types.js'; export { type VeoConfig } from './veo.js'; -async function initializer(ai: Genkit, pluginOptions?: VertexPluginOptions) { +async function initializer(pluginOptions?: VertexPluginOptions) { const clientOptions = await getDerivedOptions(pluginOptions); - veo.defineKnownModels(ai, clientOptions, pluginOptions); - imagen.defineKnownModels(ai, clientOptions, pluginOptions); - lyria.defineKnownModels(ai, clientOptions, pluginOptions); - gemini.defineKnownModels(ai, clientOptions, pluginOptions); - embedder.defineKnownModels(ai, clientOptions, pluginOptions); + return [ + ...veo.listKnownModels(clientOptions, pluginOptions), + ...imagen.listKnownModels(clientOptions, pluginOptions), + ...lyria.listKnownModels(clientOptions, pluginOptions), + ...gemini.listKnownModels(clientOptions, pluginOptions), + ...embedder.listKnownModels(clientOptions, pluginOptions), + ]; } async function resolver( - ai: Genkit, actionType: ActionType, actionName: string, pluginOptions?: VertexPluginOptions -) { +): Promise { const clientOptions = await getDerivedOptions(pluginOptions); switch (actionType) { case 'model': if (lyria.isLyriaModelName(actionName)) { - lyria.defineModel(ai, actionName, clientOptions, pluginOptions); + return lyria.defineModel(actionName, clientOptions, pluginOptions); } else if (imagen.isImagenModelName(actionName)) { - imagen.defineModel(ai, actionName, clientOptions, pluginOptions); + return imagen.defineModel(actionName, clientOptions, pluginOptions); } else if (veo.isVeoModelName(actionName)) { - // no-op (not gemini) + return undefined; } else { - gemini.defineModel(ai, actionName, clientOptions, pluginOptions); + return gemini.defineModel(actionName, clientOptions, pluginOptions); } break; case 'background-model': if (veo.isVeoModelName(actionName)) { - veo.defineModel(ai, actionName, clientOptions, pluginOptions); + return veo.defineModel(actionName, clientOptions, pluginOptions); } break; case 'embedder': - embedder.defineEmbedder(ai, actionName, clientOptions, pluginOptions); + return embedder.defineEmbedder(actionName, clientOptions, pluginOptions); break; - default: - // no-op } + return undefined; } async function listActions(options?: VertexPluginOptions) { @@ -102,23 +106,23 @@ async function listActions(options?: VertexPluginOptions) { /** * Add Google Cloud Vertex AI to Genkit. Includes Gemini and Imagen models and text embedder. */ -function vertexAIPlugin(options?: VertexPluginOptions): GenkitPlugin { +function vertexAIPlugin(options?: VertexPluginOptions): GenkitPluginV2 { let listActionsCache; - return genkitPlugin( - 'vertexai', - async (ai: Genkit) => await initializer(ai, options), - async (ai: Genkit, actionType: ActionType, actionName: string) => - await resolver(ai, actionType, actionName, options), - async () => { + return genkitPluginV2({ + name: 'vertexai', + init: async () => await initializer(options), + resolve: async (actionType: ActionType, actionName: string) => + await resolver(actionType, actionName, options), + list: async () => { if (listActionsCache) return listActionsCache; listActionsCache = await listActions(options); return listActionsCache; - } - ); + }, + }); } export type VertexAIPlugin = { - (pluginOptions?: VertexPluginOptions): GenkitPlugin; + (pluginOptions?: VertexPluginOptions): GenkitPluginV2; model( name: gemini.KnownModels | (gemini.GeminiModelName & {}), config?: gemini.GeminiConfig diff --git a/js/plugins/google-genai/src/vertexai/lyria.ts b/js/plugins/google-genai/src/vertexai/lyria.ts index 522d485c48..584e4e859f 100644 --- a/js/plugins/google-genai/src/vertexai/lyria.ts +++ b/js/plugins/google-genai/src/vertexai/lyria.ts @@ -16,13 +16,13 @@ import { ActionMetadata, - Genkit, modelActionMetadata, modelRef, ModelReference, z, } from 'genkit'; import { ModelAction, ModelInfo } from 'genkit/model'; +import { model as pluginModel } from 'genkit/plugin'; import { lyriaPredict } from './client.js'; import { fromLyriaResponse, toLyriaPredictRequest } from './converters.js'; import { ClientOptions, Model, VertexPluginOptions } from './types.js'; @@ -112,27 +112,24 @@ export function listActions(models: Model[]): ActionMetadata[] { }); } -export function defineKnownModels( - ai: Genkit, +export function listKnownModels( clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineModel(ai, name, clientOptions, pluginOptions); - } + return Object.keys(KNOWN_MODELS).map((name: string) => + defineModel(name, clientOptions, pluginOptions) + ); } export function defineModel( - ai: Genkit, name: string, clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ): ModelAction { const ref = model(name); - return ai.defineModel( + return pluginModel( { - apiVersion: 'v2', name: ref.name, ...ref.info, configSchema: ref.configSchema, diff --git a/js/plugins/google-genai/src/vertexai/veo.ts b/js/plugins/google-genai/src/vertexai/veo.ts index f44589f7a5..01c177c0ad 100644 --- a/js/plugins/google-genai/src/vertexai/veo.ts +++ b/js/plugins/google-genai/src/vertexai/veo.ts @@ -20,9 +20,9 @@ import { modelActionMetadata, modelRef, z, - type Genkit, } from 'genkit'; import { BackgroundModelAction, ModelInfo } from 'genkit/model'; +import { backgroundModel as pluginBackgroundModel } from 'genkit/plugin'; import { veoCheckOperation, veoPredict } from './client.js'; import { fromVeoOperation, @@ -171,25 +171,23 @@ export function listActions(models: Model[]): ActionMetadata[] { }); } -export function defineKnownModels( - ai: Genkit, +export function listKnownModels( clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ) { - for (const name of Object.keys(KNOWN_MODELS)) { - defineModel(ai, name, clientOptions, pluginOptions); - } + return Object.keys(KNOWN_MODELS).map((name: string) => + defineModel(name, clientOptions, pluginOptions) + ); } export function defineModel( - ai: Genkit, name: string, clientOptions: ClientOptions, pluginOptions?: VertexPluginOptions ): BackgroundModelAction { const ref = model(name); - return ai.defineBackgroundModel({ + return pluginBackgroundModel({ name: ref.name, ...ref.info, configSchema: ref.configSchema, diff --git a/js/plugins/google-genai/tests/googleai/embedder_test.ts b/js/plugins/google-genai/tests/googleai/embedder_test.ts index e1f4acc51e..944c35b254 100644 --- a/js/plugins/google-genai/tests/googleai/embedder_test.ts +++ b/js/plugins/google-genai/tests/googleai/embedder_test.ts @@ -15,7 +15,7 @@ */ import * as assert from 'assert'; -import { Document, Genkit, GenkitError } from 'genkit'; +import { Document, GenkitError } from 'genkit'; import { afterEach, beforeEach, describe, it } from 'node:test'; import * as sinon from 'sinon'; import { @@ -29,26 +29,12 @@ import { import { MISSING_API_KEY_ERROR } from '../../src/googleai/utils.js'; describe('defineGoogleAIEmbedder', () => { - let mockGenkit: sinon.SinonStubbedInstance; let fetchStub: sinon.SinonStub; const ORIGINAL_ENV = process.env; - let embedderFunc: ( - input: Document[], - options?: EmbeddingConfig - ) => Promise; - beforeEach(() => { process.env = { ...ORIGINAL_ENV }; // Shallow clone ORIGINAL_ENV - mockGenkit = sinon.createStubInstance(Genkit); fetchStub = sinon.stub(global, 'fetch'); - - mockGenkit.defineEmbedder.callsFake((config, func) => { - embedderFunc = func; - return { - name: config.name, - } as any; - }); }); afterEach(() => { @@ -69,28 +55,6 @@ describe('defineGoogleAIEmbedder', () => { apiKey: 'test-api-key-option', }; - it('defines an embedder with the correct name and info for known model', () => { - defineEmbedder(mockGenkit, 'text-embedding-004', defaultPluginOptions); - sinon.assert.calledOnce(mockGenkit.defineEmbedder); - const args = mockGenkit.defineEmbedder.lastCall.args[0]; - assert.strictEqual(args.name, 'googleai/text-embedding-004'); - assert.strictEqual(args.info?.dimensions, 768); - }); - - it('defines an embedder with a custom name', () => { - defineEmbedder(mockGenkit, 'custom-embedding-model', defaultPluginOptions); - sinon.assert.calledOnce(mockGenkit.defineEmbedder); - const args = mockGenkit.defineEmbedder.lastCall.args[0]; - assert.strictEqual(args.name, 'googleai/custom-embedding-model'); - }); - - it('handles custom name with prefix', () => { - defineEmbedder(mockGenkit, 'googleai/custom-model', defaultPluginOptions); - sinon.assert.calledOnce(mockGenkit.defineEmbedder); - const args = mockGenkit.defineEmbedder.lastCall.args[0]; - assert.strictEqual(args.name, 'googleai/custom-model'); - }); - describe('API Key Handling', () => { beforeEach(() => { // Clear potentially relevant env variables @@ -101,14 +65,19 @@ describe('defineGoogleAIEmbedder', () => { it('throws if no API key is provided in options or env', () => { assert.throws(() => { - defineEmbedder(mockGenkit, 'text-embedding-004', {}); + defineEmbedder('text-embedding-004', {}); }, MISSING_API_KEY_ERROR); }); it('uses API key from pluginOptions if provided', async () => { - defineEmbedder(mockGenkit, 'text-embedding-004', defaultPluginOptions); + const embedder = defineEmbedder( + 'text-embedding-004', + defaultPluginOptions + ); mockFetchResponse({ embedding: { values: [] } }); - await embedderFunc([new Document({ content: [{ text: 'test' }] })]); + await embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual( @@ -119,9 +88,11 @@ describe('defineGoogleAIEmbedder', () => { it('uses API key from GEMINI_API_KEY env var', async () => { process.env.GEMINI_API_KEY = 'gemini-key'; - defineEmbedder(mockGenkit, 'text-embedding-004', {}); + const embedder = defineEmbedder('text-embedding-004', {}); mockFetchResponse({ embedding: { values: [] } }); - await embedderFunc([new Document({ content: [{ text: 'test' }] })]); + await embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual(fetchOptions.headers['x-goog-api-key'], 'gemini-key'); @@ -129,9 +100,11 @@ describe('defineGoogleAIEmbedder', () => { it('uses API key from GOOGLE_API_KEY env var', async () => { process.env.GOOGLE_API_KEY = 'google-key'; - defineEmbedder(mockGenkit, 'text-embedding-004', {}); + const embedder = defineEmbedder('text-embedding-004', {}); mockFetchResponse({ embedding: { values: [] } }); - await embedderFunc([new Document({ content: [{ text: 'test' }] })]); + await embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual(fetchOptions.headers['x-goog-api-key'], 'google-key'); @@ -139,9 +112,11 @@ describe('defineGoogleAIEmbedder', () => { it('uses API key from GOOGLE_GENAI_API_KEY env var', async () => { process.env.GOOGLE_GENAI_API_KEY = 'google-genai-key'; - defineEmbedder(mockGenkit, 'text-embedding-004', {}); + const embedder = defineEmbedder('text-embedding-004', {}); mockFetchResponse({ embedding: { values: [] } }); - await embedderFunc([new Document({ content: [{ text: 'test' }] })]); + await embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual( @@ -152,9 +127,14 @@ describe('defineGoogleAIEmbedder', () => { it('pluginOptions apiKey takes precedence over env vars', async () => { process.env.GEMINI_API_KEY = 'gemini-key'; - defineEmbedder(mockGenkit, 'text-embedding-004', defaultPluginOptions); + const embedder = defineEmbedder( + 'text-embedding-004', + defaultPluginOptions + ); mockFetchResponse({ embedding: { values: [] } }); - await embedderFunc([new Document({ content: [{ text: 'test' }] })]); + await embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual( @@ -164,11 +144,13 @@ describe('defineGoogleAIEmbedder', () => { }); it('throws if apiKey is false in pluginOptions and not provided in call options', async () => { - defineEmbedder(mockGenkit, 'text-embedding-004', { + const embedder = defineEmbedder('text-embedding-004', { apiKey: false, }); await assert.rejects( - embedderFunc([new Document({ content: [{ text: 'test' }] })]), + embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + }), (err: GenkitError) => { assert.strictEqual(err.status, 'INVALID_ARGUMENT'); assert.match( @@ -182,12 +164,15 @@ describe('defineGoogleAIEmbedder', () => { }); it('uses API key from call options if apiKey is false in pluginOptions', async () => { - defineEmbedder(mockGenkit, 'text-embedding-004', { + const embedder = defineEmbedder('text-embedding-004', { apiKey: false, }); mockFetchResponse({ embedding: { values: [] } }); - await embedderFunc([new Document({ content: [{ text: 'test' }] })], { - apiKey: 'call-time-api-key', + await embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + options: { + apiKey: 'call-time-api-key', + }, }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; @@ -198,10 +183,16 @@ describe('defineGoogleAIEmbedder', () => { }); it('call options apiKey takes precedence over pluginOptions apiKey', async () => { - defineEmbedder(mockGenkit, 'text-embedding-004', defaultPluginOptions); + const embedder = defineEmbedder( + 'text-embedding-004', + defaultPluginOptions + ); mockFetchResponse({ embedding: { values: [] } }); - await embedderFunc([new Document({ content: [{ text: 'test' }] })], { - apiKey: 'call-time-api-key', + await embedder.run({ + input: [new Document({ content: [{ text: 'test' }] })], + options: { + apiKey: 'call-time-api-key', + }, }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; @@ -217,7 +208,10 @@ describe('defineGoogleAIEmbedder', () => { const testDoc2 = new Document({ content: [{ text: 'World' }] }); it('calls embedContent for each document', async () => { - defineEmbedder(mockGenkit, 'text-embedding-004', defaultPluginOptions); + const embedder = defineEmbedder( + 'text-embedding-004', + defaultPluginOptions + ); const mockResponse1: EmbedContentResponse = { embedding: { values: [0.1, 0.2] }, @@ -239,7 +233,7 @@ describe('defineGoogleAIEmbedder', () => { }) ); - const result = await embedderFunc([testDoc1, testDoc2]); + const result = await embedder.run({ input: [testDoc1, testDoc2] }); sinon.assert.calledTwice(fetchStub); const expectedUrl = @@ -261,13 +255,16 @@ describe('defineGoogleAIEmbedder', () => { }; assert.deepStrictEqual(JSON.parse(fetchArgs2[1].body), expectedRequest2); - assert.deepStrictEqual(result, { + assert.deepStrictEqual(result.result, { embeddings: [{ embedding: [0.1, 0.2] }, { embedding: [0.3, 0.4] }], }); }); it('calls embedContent with taskType, title, and outputDimensionality options', async () => { - defineEmbedder(mockGenkit, 'text-embedding-004', defaultPluginOptions); + const embedder = defineEmbedder( + 'text-embedding-004', + defaultPluginOptions + ); mockFetchResponse({ embedding: { values: [0.1] } }); const config: EmbeddingConfig = { @@ -275,7 +272,7 @@ describe('defineGoogleAIEmbedder', () => { title: 'Doc Title', outputDimensionality: 256, }; - await embedderFunc([testDoc1], config); + await embedder.run({ input: [testDoc1], options: config }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; @@ -291,10 +288,10 @@ describe('defineGoogleAIEmbedder', () => { }); it('uses the correct model name in the URL', async () => { - defineEmbedder(mockGenkit, 'custom-model', defaultPluginOptions); + const embedder = defineEmbedder('custom-model', defaultPluginOptions); mockFetchResponse({ embedding: { values: [0.1] } }); - await embedderFunc([testDoc1]); + await embedder.run({ input: [testDoc1] }); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; @@ -304,10 +301,13 @@ describe('defineGoogleAIEmbedder', () => { }); it('uses the correct model name in the URL with prefix', async () => { - defineEmbedder(mockGenkit, 'googleai/custom-model', defaultPluginOptions); + const embedder = defineEmbedder( + 'googleai/custom-model', + defaultPluginOptions + ); mockFetchResponse({ embedding: { values: [0.1] } }); - await embedderFunc([testDoc1]); + await embedder.run({ input: [testDoc1] }); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; diff --git a/js/plugins/google-genai/tests/googleai/gemini_test.ts b/js/plugins/google-genai/tests/googleai/gemini_test.ts index 82f1f4cecc..7ebb15719c 100644 --- a/js/plugins/google-genai/tests/googleai/gemini_test.ts +++ b/js/plugins/google-genai/tests/googleai/gemini_test.ts @@ -15,9 +15,8 @@ */ import * as assert from 'assert'; -import { Genkit, z } from 'genkit'; +import { z } from 'genkit'; import { GenerateRequest } from 'genkit/model'; -import { AsyncLocalStorage } from 'node:async_hooks'; import { afterEach, beforeEach, describe, it } from 'node:test'; import * as sinon from 'sinon'; import { @@ -35,18 +34,8 @@ import { import { MISSING_API_KEY_ERROR } from '../../src/googleai/utils.js'; describe('Google AI Gemini', () => { - let mockGenkit: sinon.SinonStubbedInstance; const ORIGINAL_ENV = { ...process.env }; - let modelActionCallback: ( - request: GenerateRequest, - options: { - streamingRequested?: boolean; - sendChunk?: (chunk: any) => void; - abortSignal?: AbortSignal; - } - ) => Promise; - let fetchStub: sinon.SinonStub; beforeEach(() => { @@ -55,28 +44,7 @@ describe('Google AI Gemini', () => { delete process.env.GOOGLE_API_KEY; delete process.env.GOOGLE_GENAI_API_KEY; - mockGenkit = sinon.createStubInstance(Genkit); - - // Setup mock registry and asyncStore - const mockAsyncStore = sinon.createStubInstance(AsyncLocalStorage); - mockAsyncStore.getStore.returns(undefined); // Simulate no parent span - mockAsyncStore.run.callsFake((key, store, callback) => callback()); - - (mockGenkit as any).registry = { - lookupAction: () => undefined, - lookupFlow: () => undefined, - generateTraceId: () => 'test-trace-id', - asyncStore: mockAsyncStore, - }; - fetchStub = sinon.stub(global, 'fetch'); - - mockGenkit.defineModel.callsFake((config: any, callback: any) => { - modelActionCallback = callback; - return { - name: config.name, - } as any; - }); }); afterEach(() => { @@ -133,33 +101,20 @@ describe('Google AI Gemini', () => { candidates: [mockCandidate], }; - describe('defineGeminiModel', () => { - it('defines a model with the correct name for known model', () => { - defineModel(mockGenkit, 'gemini-2.0-flash', defaultPluginOptions); - sinon.assert.calledOnce(mockGenkit.defineModel); - const args = mockGenkit.defineModel.lastCall.args[0]; - assert.strictEqual(args.name, 'googleai/gemini-2.0-flash'); - }); - - it('defines a model with a custom name', () => { - defineModel(mockGenkit, 'my-custom-gemini', defaultPluginOptions); - const args = mockGenkit.defineModel.lastCall.args[0]; - assert.strictEqual(args.name, 'googleai/my-custom-gemini'); - }); - + describe('defineModel', () => { describe('API Key Handling', () => { it('throws if no API key is provided', () => { assert.throws(() => { - defineModel(mockGenkit, 'gemini-2.0-flash'); + defineModel('gemini-2.0-flash'); }, MISSING_API_KEY_ERROR); }); it('uses API key from pluginOptions', async () => { - defineModel(mockGenkit, 'gemini-2.0-flash', { + const model = defineModel('gemini-2.0-flash', { apiKey: 'plugin-key', }); mockFetchResponse(defaultApiResponse); - await modelActionCallback(minimalRequest, {}); + await model.run(minimalRequest); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual( @@ -170,9 +125,9 @@ describe('Google AI Gemini', () => { it('uses API key from GEMINI_API_KEY env var', async () => { process.env.GEMINI_API_KEY = 'gemini-key'; - defineModel(mockGenkit, 'gemini-2.0-flash'); + const model = defineModel('gemini-2.0-flash'); mockFetchResponse(defaultApiResponse); - await modelActionCallback(minimalRequest, {}); + await model.run(minimalRequest); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual( fetchOptions.headers['x-goog-api-key'], @@ -181,22 +136,22 @@ describe('Google AI Gemini', () => { }); it('throws if apiKey is false and not in call config', async () => { - defineModel(mockGenkit, 'gemini-2.0-flash', { apiKey: false }); + const model = defineModel('gemini-2.0-flash', { apiKey: false }); await assert.rejects( - modelActionCallback(minimalRequest, {}), + model.run(minimalRequest), /GoogleAI plugin was initialized with \{apiKey: false\}/ ); sinon.assert.notCalled(fetchStub); }); it('uses API key from call config if apiKey is false', async () => { - defineModel(mockGenkit, 'gemini-2.0-flash', { apiKey: false }); + const model = defineModel('gemini-2.0-flash', { apiKey: false }); mockFetchResponse(defaultApiResponse); const request: GenerateRequest = { ...minimalRequest, config: { apiKey: 'call-time-key' }, }; - await modelActionCallback(request, {}); + await model.run(request); const fetchOptions = fetchStub.lastCall.args[1]; assert.strictEqual( fetchOptions.headers['x-goog-api-key'], @@ -206,15 +161,10 @@ describe('Google AI Gemini', () => { }); describe('Request Formation and API Calls', () => { - beforeEach(() => { - defineModel(mockGenkit, 'gemini-2.5-flash', defaultPluginOptions); - }); - it('calls fetch for non-streaming requests', async () => { + const model = defineModel('gemini-2.5-flash', defaultPluginOptions); mockFetchResponse(defaultApiResponse); - await modelActionCallback(minimalRequest, { - streamingRequested: false, - }); + await model.run(minimalRequest); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; @@ -234,13 +184,11 @@ describe('Google AI Gemini', () => { }); it('calls fetch for streaming requests', async () => { + const model = defineModel('gemini-2.5-flash', defaultPluginOptions); mockFetchStreamResponse([defaultApiResponse]); const sendChunkSpy = sinon.spy(); - await modelActionCallback(minimalRequest, { - streamingRequested: true, - sendChunk: sendChunkSpy, - }); + await model.run(minimalRequest, { onChunk: sendChunkSpy }); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; @@ -261,11 +209,11 @@ describe('Google AI Gemini', () => { }); it('passes AbortSignal to fetch', async () => { + const model = defineModel('gemini-2.5-flash', defaultPluginOptions); mockFetchResponse(defaultApiResponse); const controller = new AbortController(); const abortSignal = controller.signal; - await modelActionCallback(minimalRequest, { - streamingRequested: false, + await model.run(minimalRequest, { abortSignal, }); sinon.assert.calledOnce(fetchStub); @@ -285,6 +233,7 @@ describe('Google AI Gemini', () => { }); it('handles system instructions', async () => { + const model = defineModel('gemini-2.5-flash', defaultPluginOptions); mockFetchResponse(defaultApiResponse); const request: GenerateRequest = { messages: [ @@ -292,7 +241,7 @@ describe('Google AI Gemini', () => { { role: 'user', content: [{ text: 'Hello' }] }, ], }; - await modelActionCallback(request, {}); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body @@ -307,6 +256,7 @@ describe('Google AI Gemini', () => { }); it('constructs tools array correctly', async () => { + const model = defineModel('gemini-2.5-flash', defaultPluginOptions); mockFetchResponse(defaultApiResponse); const request: GenerateRequest = { ...minimalRequest, @@ -323,7 +273,7 @@ describe('Google AI Gemini', () => { googleSearchRetrieval: {}, }, }; - await modelActionCallback(request, {}); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body @@ -338,47 +288,42 @@ describe('Google AI Gemini', () => { }); describe('Error Handling', () => { - beforeEach(() => { - defineModel(mockGenkit, 'gemini-2.0-flash', defaultPluginOptions); - }); - it('throws if no candidates are returned', async () => { + const model = defineModel('gemini-2.0-flash', defaultPluginOptions); mockFetchResponse({ candidates: [] }); await assert.rejects( - modelActionCallback(minimalRequest, {}), + model.run(minimalRequest), /No valid candidates returned/ ); }); it('throws on fetch error', async () => { + const model = defineModel('gemini-2.0-flash', defaultPluginOptions); fetchStub.rejects(new Error('Network error')); - await assert.rejects( - modelActionCallback(minimalRequest, {}), - /Failed to fetch/ - ); + await assert.rejects(model.run(minimalRequest), /Failed to fetch/); }); }); describe('Debug Traces', () => { it('API call works with debugTraces: true', async () => { - defineModel(mockGenkit, 'gemini-2.5-flash', { + const model = defineModel('gemini-2.5-flash', { ...defaultPluginOptions, experimental_debugTraces: true, }); mockFetchResponse(defaultApiResponse); - await assert.doesNotReject(modelActionCallback(minimalRequest, {})); + await assert.doesNotReject(model.run(minimalRequest)); sinon.assert.calledOnce(fetchStub); }); it('API call works with debugTraces: false', async () => { - defineModel(mockGenkit, 'gemini-2.0-flash', { + const model = defineModel('gemini-2.0-flash', { ...defaultPluginOptions, experimental_debugTraces: false, }); mockFetchResponse(defaultApiResponse); - await assert.doesNotReject(modelActionCallback(minimalRequest, {})); + await assert.doesNotReject(model.run(minimalRequest)); sinon.assert.calledOnce(fetchStub); }); }); diff --git a/js/plugins/google-genai/tests/googleai/imagen_test.ts b/js/plugins/google-genai/tests/googleai/imagen_test.ts index 6fe76408c0..1eab8dc7ff 100644 --- a/js/plugins/google-genai/tests/googleai/imagen_test.ts +++ b/js/plugins/google-genai/tests/googleai/imagen_test.ts @@ -15,7 +15,7 @@ */ import * as assert from 'assert'; -import { Genkit, MessageData } from 'genkit'; +import { MessageData } from 'genkit'; import { GenerateRequest, getBasicUsageStats } from 'genkit/model'; import { afterEach, beforeEach, describe, it } from 'node:test'; import * as sinon from 'sinon'; @@ -159,7 +159,6 @@ describe('Google AI Imagen', () => { }); describe('defineModel()', () => { - let mockAi: sinon.SinonStubbedInstance; let fetchStub: sinon.SinonStub; let envStub: sinon.SinonStub; @@ -167,7 +166,6 @@ describe('Google AI Imagen', () => { const defaultApiKey = 'default-api-key'; beforeEach(() => { - mockAi = sinon.createStubInstance(Genkit); fetchStub = sinon.stub(global, 'fetch'); // Stub process.env to control environment variables envStub = sinon.stub(process, 'env').value({}); @@ -199,12 +197,8 @@ describe('Google AI Imagen', () => { const baseUrl = defineOptions.baseUrl; const apiKey = defineOptions.apiKey; - defineModel(mockAi as any, name, { apiKey, apiVersion, baseUrl }); - assert.ok(mockAi.defineModel.calledOnce, 'defineModel should be called'); - const callArgs = mockAi.defineModel.firstCall.args; - assert.strictEqual(callArgs[0].name, `googleai/${name}`); - assert.strictEqual(callArgs[0].configSchema, ImagenConfigSchema); - return callArgs[1]; + const model = defineModel(name, { apiKey, apiVersion, baseUrl }); + return model.run; } it('should define a model and call fetch successfully', async () => { @@ -266,13 +260,13 @@ describe('Google AI Imagen', () => { role: 'model', content: expectedContent, }; - assert.deepStrictEqual(result.message, expectedMessage); - assert.strictEqual(result.finishReason, 'stop'); + assert.deepStrictEqual(result.result.message, expectedMessage); + assert.strictEqual(result.result.finishReason, 'stop'); assert.deepStrictEqual( - result.usage, + result.result.usage, getBasicUsageStats(request.messages, expectedMessage) ); - assert.deepStrictEqual(result.custom, mockResponse); + assert.deepStrictEqual(result.result.custom, mockResponse); }); it('should use default apiKey if no request apiKey provided', async () => { @@ -420,9 +414,8 @@ describe('Google AI Imagen', () => { // process.env is empty due to envStub in beforeEach assert.throws(() => { // Explicitly pass undefined for apiKey - defineModel(mockAi as any, modelName, undefined); + defineModel(modelName, undefined); }, MISSING_API_KEY_ERROR); - sinon.assert.notCalled(mockAi.defineModel); }); it('should use key from env if no key passed to defineImagenModel', async () => { diff --git a/js/plugins/google-genai/tests/googleai/index_test.ts b/js/plugins/google-genai/tests/googleai/index_test.ts index e6ffffc6df..b70c15fba0 100644 --- a/js/plugins/google-genai/tests/googleai/index_test.ts +++ b/js/plugins/google-genai/tests/googleai/index_test.ts @@ -134,10 +134,9 @@ describe('GoogleAI Plugin', () => { delete process.env.GOOGLE_API_KEY; delete process.env.GEMINI_API_KEY; delete process.env.GOOGLE_GENAI_API_KEY; - const ai = genkit({ plugins: [googleAI({})] }); - const pluginProvider = googleAI()(ai); + const plugin = googleAI({}); await assert.rejects(async () => { - await pluginProvider.initializer(); + await plugin.init!(); }, MISSING_API_KEY_ERROR); }); @@ -145,10 +144,9 @@ describe('GoogleAI Plugin', () => { delete process.env.GOOGLE_API_KEY; delete process.env.GEMINI_API_KEY; delete process.env.GOOGLE_GENAI_API_KEY; - const ai = genkit({ plugins: [googleAI({ apiKey: false })] }); - const pluginProvider = googleAI({ apiKey: false })(ai); + const plugin = googleAI({ apiKey: false }); await assert.doesNotReject(async () => { - await pluginProvider.initializer(); + await plugin.init!(); }); }); }); @@ -325,10 +323,6 @@ describe('GoogleAI Plugin', () => { VeoConfigSchema, 'Should have VeoConfigSchema' ); - assert.ok( - modelRef.info?.supports?.longRunning, - 'Veo should support longRunning' - ); }); it('should have config values for veo model', () => { @@ -423,15 +417,13 @@ describe('GoogleAI Plugin', () => { }; it('should return an empty array if no models are returned', async () => { - const ai = genkit({ plugins: [googleAI()] }); fetchMock.mock.mockImplementation(async () => createMockResponse([])); - const pluginProvider = googleAI()(ai); - const actions = await pluginProvider.listActions!(); + 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 ai = genkit({ plugins: [googleAI()] }); const mockModels: Partial[] = [ { name: 'models/gemini-2.5-pro', @@ -458,8 +450,8 @@ describe('GoogleAI Plugin', () => { createMockResponse(mockModels) ); - const pluginProvider = googleAI()(ai); - const actions = await pluginProvider.listActions!(); + const plugin = googleAI(); + const actions = await plugin.list!(); const actionNames = actions.map((a) => a.name).sort(); assert.deepStrictEqual( actionNames, @@ -493,7 +485,6 @@ describe('GoogleAI Plugin', () => { }); it('should filter out deprecated models', async () => { - const ai = genkit({ plugins: [googleAI()] }); const mockModels = [ { name: 'models/gemini-1.5-flash', @@ -523,14 +514,13 @@ describe('GoogleAI Plugin', () => { fetchMock.mock.mockImplementation(async () => createMockResponse(mockModels) ); - const pluginProvider = googleAI()(ai); - const actions = await pluginProvider.listActions!(); + 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 () => { - const ai = genkit({ plugins: [googleAI()] }); fetchMock.mock.mockImplementation(async () => { return Promise.resolve({ ok: false, @@ -540,8 +530,8 @@ describe('GoogleAI Plugin', () => { text: async () => JSON.stringify({ error: { message: 'API Error' } }), }); }); - const pluginProvider = googleAI()(ai); - const actions = await pluginProvider.listActions!(); + const plugin = googleAI(); + const actions = await plugin.list!(); assert.deepStrictEqual( actions, [], @@ -553,10 +543,9 @@ describe('GoogleAI Plugin', () => { delete process.env.GOOGLE_API_KEY; delete process.env.GEMINI_API_KEY; delete process.env.GOOGLE_GENAI_API_KEY; - const ai = genkit({ plugins: [googleAI({ apiKey: false })] }); // Init with apiKey: false - const pluginProvider = googleAI({ apiKey: false })(ai); - const actions = await pluginProvider.listActions!(); + const plugin = googleAI({ apiKey: false }); + const actions = await plugin.list!(); assert.deepStrictEqual( actions, [], @@ -570,7 +559,6 @@ describe('GoogleAI Plugin', () => { }); it('should use listActions cache', async () => { - const ai = genkit({ plugins: [googleAI()] }); const mockModels = [ { name: 'models/gemini-1.0-pro', @@ -580,9 +568,9 @@ describe('GoogleAI Plugin', () => { fetchMock.mock.mockImplementation(async () => createMockResponse(mockModels) ); - const pluginProvider = googleAI()(ai); - await pluginProvider.listActions!(); - await pluginProvider.listActions!(); + const plugin = googleAI(); + await plugin.list!(); + await plugin.list!(); assert.strictEqual( fetchMock.mock.callCount(), 1, diff --git a/js/plugins/google-genai/tests/googleai/veo_test.ts b/js/plugins/google-genai/tests/googleai/veo_test.ts index 13ccd57364..c88e955356 100644 --- a/js/plugins/google-genai/tests/googleai/veo_test.ts +++ b/js/plugins/google-genai/tests/googleai/veo_test.ts @@ -15,7 +15,7 @@ */ import * as assert from 'assert'; -import { Genkit, Operation } from 'genkit'; +import { Operation } from 'genkit'; import { GenerateRequest, GenerateResponseData } from 'genkit/model'; import { afterEach, beforeEach, describe, it } from 'node:test'; import * as sinon from 'sinon'; @@ -27,6 +27,7 @@ import { VeoConfig, VeoConfigSchema, defineModel, + listKnownModels, model, } from '../../src/googleai/veo.js'; @@ -40,13 +41,27 @@ describe('Google AI Veo', () => { }); }); + describe('listKnownModels()', () => { + it('should return an array of model actions', () => { + const models = listKnownModels(); + assert.ok(Array.isArray(models)); + assert.strictEqual(models.length, Object.keys(KNOWN_MODELS).length); + models.forEach((m) => { + assert.ok(m.__action.name.startsWith('googleai/veo-')); + assert.ok(m.start); + assert.ok(m.check); + }); + }); + }); + describe('model()', () => { it('should return a ModelReference for a known model', () => { const modelName = 'veo-2.0-generate-001'; const ref = model(modelName); assert.strictEqual(ref.name, `googleai/${modelName}`); assert.ok(ref.info?.supports?.media); - assert.ok(ref.info?.supports?.longRunning); + // TODO: remove cast if we fix longRunning + assert.ok((ref.info?.supports as any).longRunning); }); it('should return a ModelReference for an unknown model using generic info', () => { @@ -175,7 +190,6 @@ describe('Google AI Veo', () => { }); describe('defineModel()', () => { - let mockAi: sinon.SinonStubbedInstance; let fetchStub: sinon.SinonStub; let envStub: sinon.SinonStub; @@ -183,7 +197,6 @@ describe('Google AI Veo', () => { const defaultApiKey = 'default-api-key'; beforeEach(() => { - mockAi = sinon.createStubInstance(Genkit); fetchStub = sinon.stub(global, 'fetch'); envStub = sinon.stub(process, 'env').value({}); }); @@ -219,15 +232,13 @@ describe('Google AI Veo', () => { const baseUrl = defineOptions.baseUrl; const apiKey = defineOptions.apiKey; - defineModel(mockAi as any, name, { apiKey, apiVersion, baseUrl }); - assert.ok( - mockAi.defineBackgroundModel.calledOnce, - 'defineBackgroundModel should be called' - ); - const callArgs = mockAi.defineBackgroundModel.firstCall.args; - assert.strictEqual(callArgs[0].name, `googleai/${name}`); - assert.strictEqual(callArgs[0].configSchema, VeoConfigSchema); - return { start: callArgs[0].start, check: callArgs[0].check }; + const model = defineModel(name, { apiKey, apiVersion, baseUrl }); + assert.strictEqual(model.__action.name, `googleai/${name}`); + assert.strictEqual(model.__configSchema, VeoConfigSchema); + return { + start: (req) => model.start(req), + check: (op) => model.check(op), + }; } describe('start()', () => { @@ -276,7 +287,10 @@ describe('Google AI Veo', () => { expectedVeoPredictRequest ); - assert.deepStrictEqual(result, fromVeoOperation(mockOp)); + const expectedOp = fromVeoOperation(mockOp); + assert.strictEqual(result.id, expectedOp.id); + assert.strictEqual(result.done, expectedOp.done); + assert.ok(result.action); }); it('should handle custom apiVersion and baseUrl', async () => { @@ -362,7 +376,11 @@ describe('Google AI Veo', () => { assert.deepStrictEqual(fetchArgs[1].headers, expectedHeaders); assert.strictEqual(fetchArgs[1].method, 'GET'); - assert.deepStrictEqual(result, fromVeoOperation(mockResponse)); + const expectedOp = fromVeoOperation(mockResponse); + assert.strictEqual(result.id, expectedOp.id); + assert.strictEqual(result.done, expectedOp.done); + assert.deepStrictEqual(result.output, expectedOp.output); + assert.ok(result.action); }); it('should handle custom apiVersion and baseUrl for check', async () => { diff --git a/js/plugins/google-genai/tests/vertexai/embedder_test.ts b/js/plugins/google-genai/tests/vertexai/embedder_test.ts index 98e0b3b31b..6248ceae17 100644 --- a/js/plugins/google-genai/tests/vertexai/embedder_test.ts +++ b/js/plugins/google-genai/tests/vertexai/embedder_test.ts @@ -15,7 +15,7 @@ */ import * as assert from 'assert'; -import { Document, Genkit } from 'genkit'; +import { Document } from 'genkit'; import { GoogleAuth } from 'google-auth-library'; import { afterEach, beforeEach, describe, it } from 'node:test'; import * as sinon from 'sinon'; @@ -32,7 +32,6 @@ import { } from '../../src/vertexai/types.js'; describe('defineEmbedder', () => { - let mockGenkit: sinon.SinonStubbedInstance; let fetchStub: sinon.SinonStub; let authMock: sinon.SinonStubbedInstance; @@ -51,26 +50,13 @@ describe('defineEmbedder', () => { apiKey: 'test-global-api-key', }; - let embedderFunc: ( - input: Document[], - options?: EmbeddingConfig - ) => Promise; - beforeEach(() => { - mockGenkit = sinon.createStubInstance(Genkit); fetchStub = sinon.stub(global, 'fetch'); authMock = sinon.createStubInstance(GoogleAuth); authMock.getAccessToken.resolves('test-token'); regionalClientOptions.authClient = authMock as unknown as GoogleAuth; globalClientOptions.authClient = authMock as unknown as GoogleAuth; - - mockGenkit.defineEmbedder.callsFake((config, func) => { - embedderFunc = func; - return { - name: config.name, - } as any; - }); }); afterEach(() => { @@ -105,21 +91,6 @@ describe('defineEmbedder', () => { function runTestsForClientOptions(clientOptions: ClientOptions) { describe(`with ${clientOptions.kind} client options`, () => { - it('defines an embedder with the correct name and info for known model', () => { - defineEmbedder(mockGenkit, 'text-embedding-005', clientOptions); - sinon.assert.calledOnce(mockGenkit.defineEmbedder); - const args = mockGenkit.defineEmbedder.lastCall.args[0]; - assert.strictEqual(args.name, 'vertexai/text-embedding-005'); - assert.strictEqual(args.info?.dimensions, 768); - }); - - it('defines an embedder with a custom name', () => { - defineEmbedder(mockGenkit, 'custom-model', clientOptions); - sinon.assert.calledOnce(mockGenkit.defineEmbedder); - const args = mockGenkit.defineEmbedder.lastCall.args[0]; - assert.strictEqual(args.name, 'vertexai/custom-model'); - }); - describe('Embedder Functionality', () => { const testDoc1: Document = new Document({ content: [{ text: 'Hello' }], @@ -129,7 +100,7 @@ describe('defineEmbedder', () => { }); it('calls embedContent with text-only documents', async () => { - defineEmbedder(mockGenkit, 'text-embedding-005', clientOptions); + const embedder = defineEmbedder('text-embedding-005', clientOptions); const mockResponse: EmbedContentResponse = { predictions: [ @@ -149,7 +120,7 @@ describe('defineEmbedder', () => { }; mockFetchResponse(mockResponse); - const result = await embedderFunc([testDoc1, testDoc2]); + const result = await embedder.run({ input: [testDoc1, testDoc2] }); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; @@ -174,20 +145,20 @@ describe('defineEmbedder', () => { getExpectedHeaders(clientOptions) ); - assert.deepStrictEqual(result, { + assert.deepStrictEqual(result.result, { embeddings: [{ embedding: [0.1, 0.2] }, { embedding: [0.3, 0.4] }], }); }); it('calls embedContent with taskType and title options', async () => { - defineEmbedder(mockGenkit, 'text-embedding-005', clientOptions); + const embedder = defineEmbedder('text-embedding-005', clientOptions); mockFetchResponse({ predictions: [] }); const config: EmbeddingConfig = { taskType: 'RETRIEVAL_DOCUMENT', title: 'Doc Title', }; - await embedderFunc([testDoc1], config); + await embedder.run({ input: [testDoc1], options: config }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; @@ -197,7 +168,10 @@ describe('defineEmbedder', () => { }); it('handles multimodal embeddings for images (base64)', async () => { - defineEmbedder(mockGenkit, 'multimodalembedding@001', clientOptions); + const embedder = defineEmbedder( + 'multimodalembedding@001', + clientOptions + ); const docWithImage: Document = new Document({ content: [ { text: 'A picture' }, @@ -210,7 +184,7 @@ describe('defineEmbedder', () => { }; mockFetchResponse(mockResponse); - const result = await embedderFunc([docWithImage]); + const result = await embedder.run({ input: [docWithImage] }); const expectedInstance: EmbeddingInstance = { text: 'A picture', @@ -218,11 +192,14 @@ describe('defineEmbedder', () => { }; const fetchBody = JSON.parse(fetchStub.lastCall.args[1].body); assert.deepStrictEqual(fetchBody.instances[0], expectedInstance); - assert.deepStrictEqual(result.embeddings.length, 2); + assert.deepStrictEqual(result.result.embeddings.length, 2); }); it('handles multimodal embeddings for images (gcs)', async () => { - defineEmbedder(mockGenkit, 'multimodalembedding@001', clientOptions); + const embedder = defineEmbedder( + 'multimodalembedding@001', + clientOptions + ); const docWithImage: Document = new Document({ content: [ { @@ -234,7 +211,7 @@ describe('defineEmbedder', () => { ], }); mockFetchResponse({ predictions: [] }); - await embedderFunc([docWithImage]); + await embedder.run({ input: [docWithImage] }); const expectedInstance: EmbeddingInstance = { image: { gcsUri: 'gs://bucket/image.jpg', mimeType: 'image/jpeg' }, @@ -244,11 +221,11 @@ describe('defineEmbedder', () => { }); it('passes outputDimensionality to the API call', async () => { - defineEmbedder(mockGenkit, 'text-embedding-005', clientOptions); + const embedder = defineEmbedder('text-embedding-005', clientOptions); mockFetchResponse({ predictions: [] }); const config: EmbeddingConfig = { outputDimensionality: 256 }; - await embedderFunc([testDoc1], config); + await embedder.run({ input: [testDoc1], options: config }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; @@ -268,7 +245,7 @@ describe('defineEmbedder', () => { describe('with regional client options only', () => { const clientOptions = regionalClientOptions; it('handles multimodal embeddings for video', async () => { - defineEmbedder(mockGenkit, 'multimodalembedding@001', clientOptions); + const embedder = defineEmbedder('multimodalembedding@001', clientOptions); const docWithVideo: Document = new Document({ content: [ { text: 'A video' }, @@ -296,7 +273,7 @@ describe('defineEmbedder', () => { }; mockFetchResponse(mockResponse); - const result = await embedderFunc([docWithVideo]); + const result = await embedder.run({ input: [docWithVideo] }); const expectedInstance: EmbeddingInstance = { text: 'A video', @@ -312,7 +289,7 @@ describe('defineEmbedder', () => { const fetchBody = JSON.parse(fetchStub.lastCall.args[1].body); assert.deepStrictEqual(fetchBody.instances[0], expectedInstance); - assert.deepStrictEqual(result, { + assert.deepStrictEqual(result.result, { embeddings: [ { embedding: [0.1], metadata: { embedType: 'textEmbedding' } }, { @@ -336,12 +313,12 @@ describe('defineEmbedder', () => { }); it('throws on unsupported media type', async () => { - defineEmbedder(mockGenkit, 'multimodalembedding@001', clientOptions); + const embedder = defineEmbedder('multimodalembedding@001', clientOptions); const docWithInvalidMedia: Document = new Document({ content: [{ media: { url: 'a', contentType: 'application/pdf' } }], }); await assert.rejects( - embedderFunc([docWithInvalidMedia]), + embedder.run({ input: [docWithInvalidMedia] }), /Unsupported contentType: 'application\/pdf/ ); sinon.assert.notCalled(fetchStub); diff --git a/js/plugins/google-genai/tests/vertexai/gemini_test.ts b/js/plugins/google-genai/tests/vertexai/gemini_test.ts index cba42ff9c6..19d3d0f3f7 100644 --- a/js/plugins/google-genai/tests/vertexai/gemini_test.ts +++ b/js/plugins/google-genai/tests/vertexai/gemini_test.ts @@ -15,10 +15,9 @@ */ import * as assert from 'assert'; -import { Genkit, z } from 'genkit'; +import { z } from 'genkit'; import { GenerateRequest, ModelReference } from 'genkit/model'; import { GoogleAuth } from 'google-auth-library'; -import { AsyncLocalStorage } from 'node:async_hooks'; import { afterEach, beforeEach, describe, it } from 'node:test'; import * as sinon from 'sinon'; import { FinishReason } from '../../src/common/types.js'; @@ -40,18 +39,7 @@ import { } from '../../src/vertexai/types.js'; describe('Vertex AI Gemini', () => { - let mockGenkit: sinon.SinonStubbedInstance; - let modelActionCallback: ( - request: GenerateRequest, - options: { - streamingRequested?: boolean; - sendChunk?: (chunk: any) => void; - abortSignal?: AbortSignal; - } - ) => Promise; - let fetchStub: sinon.SinonStub; - let mockAsyncStore: sinon.SinonStubbedInstance>; let authMock: sinon.SinonStubbedInstance; const defaultRegionalClientOptions: ClientOptions = { @@ -75,30 +63,13 @@ describe('Vertex AI Gemini', () => { }; beforeEach(() => { - mockGenkit = sinon.createStubInstance(Genkit); - mockAsyncStore = sinon.createStubInstance(AsyncLocalStorage); authMock = sinon.createStubInstance(GoogleAuth); authMock.getAccessToken.resolves('test-token'); defaultRegionalClientOptions.authClient = authMock as unknown as GoogleAuth; defaultGlobalClientOptions.authClient = authMock as unknown as GoogleAuth; - mockAsyncStore.getStore.returns(undefined); - mockAsyncStore.run.callsFake((_, callback) => callback()); - - (mockGenkit as any).registry = { - lookupAction: () => undefined, - lookupFlow: () => undefined, - generateTraceId: () => 'test-trace-id', - asyncStore: mockAsyncStore, - }; - fetchStub = sinon.stub(global, 'fetch'); - - mockGenkit.defineModel.callsFake((config: any, func: any) => { - modelActionCallback = func; - return { name: config.name } as any; - }); }); afterEach(() => { @@ -178,10 +149,6 @@ describe('Vertex AI Gemini', () => { function runCommonTests(clientOptions: ClientOptions) { describe(`Model Action Callback ${clientOptions.kind}`, () => { - beforeEach(() => { - defineModel(mockGenkit, 'gemini-2.5-flash', clientOptions); - }); - function getExpectedHeaders( configApiKey?: string ): Record { @@ -233,15 +200,17 @@ describe('Vertex AI Gemini', () => { } it('throws if no messages are provided', async () => { + const model = defineModel('gemini-2.5-flash', clientOptions); await assert.rejects( - modelActionCallback({ messages: [], config: {} }, {}), + model.run({ messages: [], config: {} }), /No messages provided/ ); }); it('calls fetch for non-streaming requests', async () => { mockFetchResponse(defaultApiResponse); - const result = await modelActionCallback(minimalRequest, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + const result = await model.run(minimalRequest); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; @@ -262,9 +231,10 @@ describe('Vertex AI Gemini', () => { assert.deepStrictEqual(options.headers, getExpectedHeaders()); - assert.strictEqual(result.candidates.length, 1); + assert.ok(result.result.candidates); + assert.strictEqual(result.result.candidates.length, 1); assert.strictEqual( - result.candidates[0].message.content[0].text, + result.result.candidates[0].message.content[0].text, 'Hi there' ); }); @@ -273,10 +243,8 @@ describe('Vertex AI Gemini', () => { mockFetchStreamResponse([defaultApiResponse]); const sendChunkSpy = sinon.spy(); - await modelActionCallback(minimalRequest, { - streamingRequested: true, - sendChunk: sendChunkSpy, - }); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(minimalRequest, { onChunk: sendChunkSpy }); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; @@ -302,8 +270,8 @@ describe('Vertex AI Gemini', () => { mockFetchResponse(defaultApiResponse); const controller = new AbortController(); const abortSignal = controller.signal; - await modelActionCallback(minimalRequest, { - streamingRequested: false, + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(minimalRequest, { abortSignal, }); sinon.assert.calledOnce(fetchStub); @@ -331,7 +299,8 @@ describe('Vertex AI Gemini', () => { ], config: {}, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body @@ -351,7 +320,8 @@ describe('Vertex AI Gemini', () => { ...minimalRequest, config: { temperature: 0.1, topP: 0.8, maxOutputTokens: 100 }, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body @@ -368,7 +338,8 @@ describe('Vertex AI Gemini', () => { ...minimalRequest, config: { labels: myLabels }, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body @@ -390,7 +361,8 @@ describe('Vertex AI Gemini', () => { ], config: {}, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body @@ -417,7 +389,8 @@ describe('Vertex AI Gemini', () => { googleSearchRetrieval: {}, }, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body ); @@ -441,7 +414,8 @@ describe('Vertex AI Gemini', () => { }, }, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body ); @@ -475,7 +449,8 @@ describe('Vertex AI Gemini', () => { ], }, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body ); @@ -494,7 +469,8 @@ describe('Vertex AI Gemini', () => { output: { format: 'json' }, config: {}, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body ); @@ -506,16 +482,18 @@ describe('Vertex AI Gemini', () => { it('throws if no candidates are returned', async () => { mockFetchResponse({ candidates: [] }); + const model = defineModel('gemini-2.5-flash', clientOptions); await assert.rejects( - modelActionCallback(minimalRequest, {}), + model.run(minimalRequest), /No valid candidates returned/ ); }); it('handles API call error', async () => { mockFetchResponse({ error: { message: 'API Error' } }, 400); + const model = defineModel('gemini-2.5-flash', clientOptions); await assert.rejects( - modelActionCallback(minimalRequest, {}), + model.run(minimalRequest), /Error fetching from .*?: \[400 Error\] API Error/ ); }); @@ -527,7 +505,8 @@ describe('Vertex AI Gemini', () => { ...minimalRequest, config: { apiKey: overrideKey }, }; - await modelActionCallback(request, {}); + const model = defineModel('gemini-2.5-flash', clientOptions); + await model.run(request); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; const url = fetchArgs[0]; @@ -547,17 +526,10 @@ describe('Vertex AI Gemini', () => { } describe('defineModel - Regional Client', () => { - it('defines a model with the correct name', () => { - defineModel(mockGenkit, 'gemini-2.0-flash', defaultRegionalClientOptions); - sinon.assert.calledOnce(mockGenkit.defineModel); - const args = mockGenkit.defineModel.lastCall.args[0]; - assert.strictEqual(args.name, 'vertexai/gemini-2.0-flash'); - }); - runCommonTests(defaultRegionalClientOptions); it('handles googleSearchRetrieval tool for gemini-1.5', async () => { - defineModel(mockGenkit, 'gemini-1.5-pro', defaultRegionalClientOptions); + const model = defineModel('gemini-1.5-pro', defaultRegionalClientOptions); mockFetchResponse(defaultApiResponse); const request: GenerateRequest = { ...minimalRequest, @@ -565,7 +537,7 @@ describe('Vertex AI Gemini', () => { googleSearchRetrieval: {}, }, }; - await modelActionCallback(request, {}); + await model.run(request); const apiRequest: GenerateContentRequest = JSON.parse( fetchStub.lastCall.args[1].body ); @@ -582,24 +554,10 @@ describe('Vertex AI Gemini', () => { }); describe('defineModel - Global Client', () => { - it('defines a model with the correct name', () => { - defineModel(mockGenkit, 'gemini-2.0-flash', defaultGlobalClientOptions); - sinon.assert.calledOnce(mockGenkit.defineModel); - const args = mockGenkit.defineModel.lastCall.args[0]; - assert.strictEqual(args.name, 'vertexai/gemini-2.0-flash'); - }); - runCommonTests(defaultGlobalClientOptions); }); describe('defineModel - Express Client', () => { - it('defines a model with the correct name', () => { - defineModel(mockGenkit, 'gemini-2.0-flash', defaultExpressClientOptions); - sinon.assert.calledOnce(mockGenkit.defineModel); - const args = mockGenkit.defineModel.lastCall.args[0]; - assert.strictEqual(args.name, 'vertexai/gemini-2.0-flash'); - }); - runCommonTests(defaultExpressClientOptions); }); }); diff --git a/js/plugins/google-genai/tests/vertexai/imagen_test.ts b/js/plugins/google-genai/tests/vertexai/imagen_test.ts index 9f3359456f..7f055e5289 100644 --- a/js/plugins/google-genai/tests/vertexai/imagen_test.ts +++ b/js/plugins/google-genai/tests/vertexai/imagen_test.ts @@ -15,7 +15,6 @@ */ import * as assert from 'assert'; -import { Genkit } from 'genkit'; import { GenerateRequest, getBasicUsageStats } from 'genkit/model'; import { GoogleAuth } from 'google-auth-library'; import { afterEach, beforeEach, describe, it } from 'node:test'; @@ -91,7 +90,6 @@ describe('Vertex AI Imagen', () => { }); describe('defineImagenModel()', () => { - let mockAi: sinon.SinonStubbedInstance; let fetchStub: sinon.SinonStub; const modelName = 'imagen-test-model'; let authMock: sinon.SinonStubbedInstance; @@ -112,7 +110,6 @@ describe('Vertex AI Imagen', () => { }; beforeEach(() => { - mockAi = sinon.createStubInstance(Genkit); fetchStub = sinon.stub(global, 'fetch'); authMock = sinon.createStubInstance(GoogleAuth); authMock.getAccessToken.resolves('test-token'); @@ -136,12 +133,8 @@ describe('Vertex AI Imagen', () => { function captureModelRunner( clientOptions: ClientOptions ): (request: GenerateRequest, options: any) => Promise { - defineModel(mockAi as any, modelName, clientOptions); - assert.ok(mockAi.defineModel.calledOnce); - const callArgs = mockAi.defineModel.firstCall.args; - assert.strictEqual(callArgs[0].name, `vertexai/${modelName}`); - assert.strictEqual(callArgs[0].configSchema, ImagenConfigSchema); - return callArgs[1]; + const model = defineModel(modelName, clientOptions); + return model.run; } function getExpectedHeaders( @@ -208,12 +201,12 @@ describe('Vertex AI Imagen', () => { const expectedResponse = fromImagenResponse(mockResponse, request); const expectedCandidates = expectedResponse.candidates; - assert.deepStrictEqual(result.candidates, expectedCandidates); - assert.deepStrictEqual(result.usage, { + assert.deepStrictEqual(result.result.candidates, expectedCandidates); + assert.deepStrictEqual(result.result.usage, { ...getBasicUsageStats(request.messages, expectedCandidates as any), custom: { generations: 2 }, }); - assert.deepStrictEqual(result.custom, mockResponse); + assert.deepStrictEqual(result.result.custom, mockResponse); }); it(`should throw an error if model returns no predictions for ${clientOptions.kind}`, async () => { diff --git a/js/plugins/google-genai/tests/vertexai/index_test.ts b/js/plugins/google-genai/tests/vertexai/index_test.ts index 1a3a23c190..2e11e89f44 100644 --- a/js/plugins/google-genai/tests/vertexai/index_test.ts +++ b/js/plugins/google-genai/tests/vertexai/index_test.ts @@ -15,7 +15,7 @@ */ import * as assert from 'assert'; -import { genkit, type Genkit } from 'genkit'; +import { genkit } from 'genkit'; import { GenerateRequest } from 'genkit/model'; import { GoogleAuth } from 'google-auth-library'; import { afterEach, beforeEach, describe, it, mock } from 'node:test'; @@ -67,7 +67,7 @@ describe('VertexAI Plugin', () => { message: NOT_SUPPORTED_IN_EXPRESS_ERROR.message, }; - let ai: Genkit; + let ai: any; // Default to regional options for most tests beforeEach(() => { @@ -252,8 +252,8 @@ describe('VertexAI Plugin', () => { it('should return an empty array if no models are returned', async () => { fetchMock.mock.mockImplementation(async () => createMockResponse([])); - const pluginProvider = vertexAI()(ai); - const actions = await pluginProvider.listActions!(); + const plugin = vertexAI(); + const actions = await plugin.list!(); assert.deepStrictEqual(actions, [], 'Should return an empty array'); }); @@ -267,8 +267,8 @@ describe('VertexAI Plugin', () => { fetchMock.mock.mockImplementation(async () => createMockResponse(mockModels) ); - const pluginProvider = vertexAI()(ai); - const actions = await pluginProvider.listActions!(); + const plugin = vertexAI(); + const actions = await plugin.list!(); const actionNames = actions.map((a) => a.name).sort(); assert.deepStrictEqual( actionNames, @@ -281,8 +281,8 @@ describe('VertexAI Plugin', () => { it('should call fetch with auth token and location-specific URL for local options', async () => { fetchMock.mock.mockImplementation(async () => createMockResponse([])); - const pluginProvider = vertexAI()(ai); - await pluginProvider.listActions!(); + const plugin = vertexAI(); + await plugin.list!(); const fetchCall = fetchMock.mock.calls[0]; const headers = fetchCall.arguments[1].headers; @@ -303,8 +303,8 @@ describe('VertexAI Plugin', () => { UTILS_TEST_ONLY.setMockDerivedOptions(globalWithOptions); ai = genkit({ plugins: [vertexAI()] }); // Re-init fetchMock.mock.mockImplementation(async () => createMockResponse([])); - const pluginProvider = vertexAI()(ai); - await pluginProvider.listActions!(); + const plugin = vertexAI(); + await plugin.list!(); const fetchCall = fetchMock.mock.calls[0]; const headers = fetchCall.arguments[1].headers; @@ -322,8 +322,8 @@ describe('VertexAI Plugin', () => { UTILS_TEST_ONLY.setMockDerivedOptions(expressMockDerivedOptions); ai = genkit({ plugins: [vertexAI()] }); // Re-init fetchMock.mock.mockImplementation(async () => createMockResponse([])); - const pluginProvider = vertexAI()(ai); - const actions = await pluginProvider.listActions!(); + const plugin = vertexAI(); + const actions = await plugin.list!(); assert.strictEqual(actions.length, 0); assert.strictEqual(fetchMock.mock.calls.length, 0); }); diff --git a/js/plugins/google-genai/tests/vertexai/lyria_test.ts b/js/plugins/google-genai/tests/vertexai/lyria_test.ts index cf2698ab04..af11cd59fa 100644 --- a/js/plugins/google-genai/tests/vertexai/lyria_test.ts +++ b/js/plugins/google-genai/tests/vertexai/lyria_test.ts @@ -15,7 +15,6 @@ */ import * as assert from 'assert'; -import { Genkit } from 'genkit'; import { GenerateRequest } from 'genkit/model'; import { GoogleAuth } from 'google-auth-library'; import { afterEach, beforeEach, describe, it } from 'node:test'; @@ -40,15 +39,8 @@ import { const { GENERIC_MODEL, KNOWN_MODELS } = TEST_ONLY; describe('Vertex AI Lyria', () => { - let mockGenkit: sinon.SinonStubbedInstance; let fetchStub: sinon.SinonStub; let authMock: sinon.SinonStubbedInstance; - let modelActionCallback: ( - request: GenerateRequest, - options: { - abortSignal?: AbortSignal; - } - ) => Promise; const modelName = 'lyria-test-model'; @@ -60,17 +52,11 @@ describe('Vertex AI Lyria', () => { }; beforeEach(() => { - mockGenkit = sinon.createStubInstance(Genkit); fetchStub = sinon.stub(global, 'fetch'); authMock = sinon.createStubInstance(GoogleAuth); authMock.getAccessToken.resolves('test-token'); defaultRegionalClientOptions.authClient = authMock as unknown as GoogleAuth; - - mockGenkit.defineModel.callsFake((config: any, func: any) => { - modelActionCallback = func; - return { name: config.name } as any; - }); }); afterEach(() => { @@ -122,13 +108,6 @@ describe('Vertex AI Lyria', () => { }); describe('defineModel()', () => { - beforeEach(() => { - defineModel(mockGenkit, modelName, defaultRegionalClientOptions); - sinon.assert.calledOnce(mockGenkit.defineModel); - const args = mockGenkit.defineModel.lastCall.args[0]; - assert.strictEqual(args.name, `vertexai/${modelName}`); - }); - const prompt = 'A funky bass line'; const minimalRequest: GenerateRequest = { messages: [{ role: 'user', content: [{ text: prompt }] }], @@ -151,7 +130,8 @@ describe('Vertex AI Lyria', () => { it('should call fetch with correct params and return lyria response', async () => { mockFetchResponse(mockPrediction); - const result = await modelActionCallback(minimalRequest, {}); + const model = defineModel(modelName, defaultRegionalClientOptions); + const result = await model.run(minimalRequest); sinon.assert.calledOnce(fetchStub); const fetchArgs = fetchStub.lastCall.args; @@ -175,18 +155,24 @@ describe('Vertex AI Lyria', () => { mockPrediction, minimalRequest ); - assert.deepStrictEqual(result, expectedResponse); - assert.strictEqual(result.candidates?.length, 2); + assert.deepStrictEqual( + result.result.candidates, + expectedResponse.candidates + ); + assert.deepStrictEqual(result.result.usage, expectedResponse.usage); + assert.deepStrictEqual(result.result.custom, expectedResponse.custom); + assert.strictEqual(result.result.candidates?.length, 2); assert.strictEqual( - result.candidates[0].message.content[0].media?.url, + result.result.candidates[0].message.content[0].media?.url, 'data:audio/wav;base64,base64audio1' ); }); it('should throw if no predictions are returned', async () => { mockFetchResponse({ predictions: [] }); + const model = defineModel(modelName, defaultRegionalClientOptions); await assert.rejects( - modelActionCallback(minimalRequest, {}), + model.run(minimalRequest), /Model returned no predictions/ ); }); @@ -195,8 +181,9 @@ describe('Vertex AI Lyria', () => { const errorBody = { error: { message: 'Quota exceeded', code: 429 } }; mockFetchResponse(errorBody, 429); + const model = defineModel(modelName, defaultRegionalClientOptions); await assert.rejects( - modelActionCallback(minimalRequest, {}), + model.run(minimalRequest), /Error fetching from .*predict.* Quota exceeded/ ); }); @@ -206,14 +193,13 @@ describe('Vertex AI Lyria', () => { const controller = new AbortController(); const abortSignal = controller.signal; - // We need to re-register to pass the clientOptions with the signal const clientOptionsWithSignal = { ...defaultRegionalClientOptions, signal: abortSignal, }; - defineModel(mockGenkit, modelName, clientOptionsWithSignal); + const model = defineModel(modelName, clientOptionsWithSignal); - await modelActionCallback(minimalRequest, { abortSignal }); + await model.run(minimalRequest, { abortSignal }); sinon.assert.calledOnce(fetchStub); const fetchOptions = fetchStub.lastCall.args[1]; diff --git a/js/plugins/google-genai/tests/vertexai/veo_test.ts b/js/plugins/google-genai/tests/vertexai/veo_test.ts index 8c6dd10f4c..eea338c242 100644 --- a/js/plugins/google-genai/tests/vertexai/veo_test.ts +++ b/js/plugins/google-genai/tests/vertexai/veo_test.ts @@ -15,7 +15,7 @@ */ import * as assert from 'assert'; -import { Genkit, Operation } from 'genkit'; +import { Operation } from 'genkit'; import { GenerateRequest } from 'genkit/model'; import { GoogleAuth } from 'google-auth-library'; import { afterEach, beforeEach, describe, it } from 'node:test'; @@ -44,7 +44,6 @@ import { const { GENERIC_MODEL, KNOWN_MODELS } = TEST_ONLY; describe('Vertex AI Veo', () => { - let mockGenkit: sinon.SinonStubbedInstance; let fetchStub: sinon.SinonStub; let authMock: sinon.SinonStubbedInstance; @@ -58,18 +57,11 @@ describe('Vertex AI Veo', () => { }; beforeEach(() => { - mockGenkit = sinon.createStubInstance(Genkit); fetchStub = sinon.stub(global, 'fetch'); authMock = sinon.createStubInstance(GoogleAuth); authMock.getAccessToken.resolves('test-token'); defaultRegionalClientOptions.authClient = authMock as unknown as GoogleAuth; - - // Mock Genkit registry methods if needed, though defineBackgroundModel is the key - (mockGenkit as any).registry = { - lookupAction: () => undefined, - generateTraceId: () => 'test-trace-id', - }; }); afterEach(() => { @@ -101,7 +93,6 @@ describe('Vertex AI Veo', () => { const ref = model(knownModelName); assert.strictEqual(ref.name, `vertexai/${knownModelName}`); assert.ok(ref.info?.supports?.media); - assert.ok(ref.info?.supports?.longRunning); }); it('should return a ModelReference for an unknown model using generic info', () => { @@ -119,14 +110,10 @@ describe('Vertex AI Veo', () => { ) => Promise; check: (operation: Operation) => Promise; } { - defineModel(mockGenkit, modelName, clientOptions); - sinon.assert.calledOnce(mockGenkit.defineBackgroundModel); - const callArgs = mockGenkit.defineBackgroundModel.firstCall.args; - assert.strictEqual(callArgs[0].name, `vertexai/${modelName}`); - assert.strictEqual(callArgs[0].configSchema, VeoConfigSchema); + const model = defineModel(modelName, clientOptions); return { - start: callArgs[0].start, - check: callArgs[0].check, + start: (req) => model.start(req), + check: (op) => model.check(op), }; } @@ -172,7 +159,9 @@ describe('Vertex AI Veo', () => { expectedPredictRequest ); - assert.deepStrictEqual(result, fromVeoOperation(mockOp)); + const expectedOp = fromVeoOperation(mockOp); + assert.strictEqual(result.id, expectedOp.id); + assert.strictEqual(result.done, expectedOp.done); }); it('should propagate API errors', async () => { @@ -253,7 +242,10 @@ describe('Vertex AI Veo', () => { toVeoOperationRequest(pendingOp); assert.deepStrictEqual(JSON.parse(options.body), expectedCheckRequest); - assert.deepStrictEqual(result, fromVeoOperation(mockResponse)); + const expectedOp = fromVeoOperation(mockResponse); + assert.strictEqual(result.id, expectedOp.id); + assert.strictEqual(result.done, expectedOp.done); + assert.deepStrictEqual(result.output, expectedOp.output); }); it('should propagate API errors for check', async () => { diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index 2fa5555c98..d392a9d21b 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -181,7 +181,7 @@ importers: optionalDependencies: '@genkit-ai/firebase': specifier: ^1.16.1 - version: 1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) + version: 1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) doc-snippets: dependencies: @@ -987,19 +987,25 @@ importers: '@genkit-ai/google-cloud': specifier: workspace:* version: link:../../plugins/google-cloud - '@genkit-ai/googleai': - specifier: workspace:* - version: link:../../plugins/googleai - '@genkit-ai/vertexai': + '@genkit-ai/google-genai': specifier: workspace:* - version: link:../../plugins/vertexai + version: link:../../plugins/google-genai express: specifier: ^4.20.0 version: 4.21.2 genkit: specifier: workspace:* version: link:../../genkit + node-fetch: + specifier: 3.3.2 + version: 3.3.2 + wav: + specifier: ^1.0.2 + version: 1.0.2 devDependencies: + '@types/wav': + specifier: ^1.0.4 + version: 1.0.4 typescript: specifier: ^5.6.2 version: 5.8.3 @@ -1011,7 +1017,7 @@ importers: version: link:../../plugins/compat-oai '@genkit-ai/express': specifier: ^1.1.0 - version: 1.12.0(@genkit-ai/core@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(express@5.1.0)(genkit@genkit) + version: 1.12.0(@genkit-ai/core@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(express@5.1.0)(genkit@genkit) genkit: specifier: workspace:* version: link:../../genkit @@ -1655,7 +1661,7 @@ importers: version: link:../../plugins/ollama genkitx-openai: specifier: ^0.10.1 - version: 0.10.1(@genkit-ai/ai@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(@genkit-ai/core@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(encoding@0.1.13) + version: 0.10.1(@genkit-ai/ai@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(@genkit-ai/core@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(encoding@0.1.13) devDependencies: rimraf: specifier: ^6.0.1 @@ -2656,11 +2662,11 @@ packages: '@firebase/webchannel-wrapper@1.0.3': resolution: {integrity: sha512-2xCRM9q9FlzGZCdgDMJwc0gyUkWFtkosy7Xxr6sFgQwn+wMNIWd7xIvYNauU1r64B5L5rsGKy/n9TKJ0aAFeqQ==} - '@genkit-ai/ai@1.19.1': - resolution: {integrity: sha512-ZMre8Wpq3Vw6IPGAVL3ZnHevVPVD98o33EmiyreZQrFrRxFgSX68l00IcbQyGxCw8l2ucusGjrs4CFV82Txurg==} + '@genkit-ai/ai@1.20.0-rc.2': + resolution: {integrity: sha512-XTTTNvkC/2w9gSMxKTXaNFc2UpafP83CtvkXbSL2GNja7Gtta7pDyweg2TUOJPzeD5aG62ztuVdBrJYmv0LFNA==} - '@genkit-ai/core@1.19.1': - resolution: {integrity: sha512-NQ2h+9V88MK0CRSebd91V932z4BAVlY9wRbpFkbfcimyLS+743Wbg9ufWGrRw3bmXld8vIK5mkH2u4DJb1VrEw==} + '@genkit-ai/core@1.20.0-rc.2': + resolution: {integrity: sha512-D2RikUJRtR62E2PhNGbF4bquYmyweMokMgAZZyr/7GRSkR5x3P2FoJRAvZMLpCwq8u2LzPC80KreUZSw0Jv6jQ==} '@genkit-ai/express@1.12.0': resolution: {integrity: sha512-QAxSS07dX5ovSfsUB4s90KaDnv4zg1wnoxCZCa+jBsYUyv9NvCCTsOk25xAQgGxc7xi3+MD+3AsPier5oZILIg==} @@ -5301,8 +5307,8 @@ packages: resolution: {integrity: sha512-Jh/AIwwgaxan+7ZUUmRLCjtchyDiqh4KjBJ5tW3plBZb5iL/BPcso8A5DlzeD9qlw0duCamnNdpFjxwaT0KyKg==} engines: {node: '>=14'} - genkit@1.19.1: - resolution: {integrity: sha512-Wsa4kn1Apu3A0AgMlO51yNgK8Y+/RnSGO7iE6oKmrL2LEh1a1LtVLK0X2etZBHRYVrf0iJG5U6ClqKNR4GDWrQ==} + genkit@1.20.0-rc.2: + resolution: {integrity: sha512-M+b1qsom9HJoUwcX3KS5R18AORnSjdcg2OCnSnx7gNrd3M69MX1HlGYYnrlIwlQb5chK/+cCtFEzOcLgjl53WA==} genkitx-openai@0.10.1: resolution: {integrity: sha512-E9/DzyQcBUSTy81xT2pvEmdnn9Q/cKoojEt6lD/EdOeinhqE9oa59d/kuXTokCMekTrj3Rk7LtNBQIDjnyjNOA==} @@ -8504,9 +8510,9 @@ snapshots: '@firebase/webchannel-wrapper@1.0.3': {} - '@genkit-ai/ai@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': + '@genkit-ai/ai@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': dependencies: - '@genkit-ai/core': 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) + '@genkit-ai/core': 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) '@opentelemetry/api': 1.9.0 '@types/node': 20.19.1 colorette: 2.0.20 @@ -8525,9 +8531,9 @@ snapshots: - supports-color optional: true - '@genkit-ai/ai@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': + '@genkit-ai/ai@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': dependencies: - '@genkit-ai/core': 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) + '@genkit-ai/core': 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) '@opentelemetry/api': 1.9.0 '@types/node': 20.19.1 colorette: 2.0.20 @@ -8545,7 +8551,7 @@ snapshots: - genkit - supports-color - '@genkit-ai/core@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': + '@genkit-ai/core@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': dependencies: '@opentelemetry/api': 1.9.0 '@opentelemetry/context-async-hooks': 1.25.1(@opentelemetry/api@1.9.0) @@ -8567,7 +8573,7 @@ snapshots: zod: 3.25.67 zod-to-json-schema: 3.24.5(zod@3.25.67) optionalDependencies: - '@genkit-ai/firebase': 1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) + '@genkit-ai/firebase': 1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) transitivePeerDependencies: - '@google-cloud/firestore' - encoding @@ -8577,7 +8583,7 @@ snapshots: - supports-color optional: true - '@genkit-ai/core@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': + '@genkit-ai/core@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': dependencies: '@opentelemetry/api': 1.9.0 '@opentelemetry/context-async-hooks': 1.25.1(@opentelemetry/api@1.9.0) @@ -8608,9 +8614,9 @@ snapshots: - genkit - supports-color - '@genkit-ai/express@1.12.0(@genkit-ai/core@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(express@5.1.0)(genkit@genkit)': + '@genkit-ai/express@1.12.0(@genkit-ai/core@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(express@5.1.0)(genkit@genkit)': dependencies: - '@genkit-ai/core': 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) + '@genkit-ai/core': 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) body-parser: 1.20.3 cors: 2.8.5 express: 5.1.0 @@ -8618,12 +8624,12 @@ snapshots: transitivePeerDependencies: - supports-color - '@genkit-ai/firebase@1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': + '@genkit-ai/firebase@1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': dependencies: - '@genkit-ai/google-cloud': 1.16.1(encoding@0.1.13)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) + '@genkit-ai/google-cloud': 1.16.1(encoding@0.1.13)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) '@google-cloud/firestore': 7.11.1(encoding@0.1.13) firebase-admin: 13.4.0(encoding@0.1.13) - genkit: 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1) + genkit: 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1) optionalDependencies: firebase: 11.9.1 transitivePeerDependencies: @@ -8644,7 +8650,7 @@ snapshots: - supports-color optional: true - '@genkit-ai/google-cloud@1.16.1(encoding@0.1.13)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': + '@genkit-ai/google-cloud@1.16.1(encoding@0.1.13)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1))': dependencies: '@google-cloud/logging-winston': 6.0.1(encoding@0.1.13)(winston@3.17.0) '@google-cloud/opentelemetry-cloud-monitoring-exporter': 0.19.0(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-metrics@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) @@ -8660,7 +8666,7 @@ snapshots: '@opentelemetry/sdk-metrics': 1.25.1(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-node': 0.52.1(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-trace-base': 1.25.1(@opentelemetry/api@1.9.0) - genkit: 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1) + genkit: 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1) google-auth-library: 9.15.1(encoding@0.1.13) node-fetch: 3.3.2 winston: 3.17.0 @@ -11685,10 +11691,10 @@ snapshots: - encoding - supports-color - genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1): + genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1): dependencies: - '@genkit-ai/ai': 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) - '@genkit-ai/core': 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) + '@genkit-ai/ai': 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) + '@genkit-ai/core': 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)) uuid: 10.0.0 transitivePeerDependencies: - '@google-cloud/firestore' @@ -11698,10 +11704,10 @@ snapshots: - supports-color optional: true - genkitx-openai@0.10.1(@genkit-ai/ai@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(@genkit-ai/core@1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(encoding@0.1.13): + genkitx-openai@0.10.1(@genkit-ai/ai@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(@genkit-ai/core@1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(encoding@0.1.13): dependencies: - '@genkit-ai/ai': 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) - '@genkit-ai/core': 1.19.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) + '@genkit-ai/ai': 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) + '@genkit-ai/core': 1.20.0-rc.2(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.4.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) openai: 4.104.0(encoding@0.1.13)(zod@3.25.67) zod: 3.25.67 transitivePeerDependencies: diff --git a/js/testapps/basic-gemini/my_room.png b/js/testapps/basic-gemini/my_room.png new file mode 100644 index 0000000000..4b54e0e932 Binary files /dev/null and b/js/testapps/basic-gemini/my_room.png differ diff --git a/js/testapps/basic-gemini/package.json b/js/testapps/basic-gemini/package.json index bdc61b6022..027990b8ed 100644 --- a/js/testapps/basic-gemini/package.json +++ b/js/testapps/basic-gemini/package.json @@ -7,7 +7,8 @@ "test": "echo \"Error: no test specified\" && exit 1", "start": "node lib/index.js", "build": "tsc", - "build:watch": "tsc --watch" + "build:watch": "tsc --watch", + "genkit:dev": "genkit start -- npx tsx --watch src/index.ts" }, "keywords": [], "author": "", @@ -16,11 +17,13 @@ "genkit": "workspace:*", "@genkit-ai/firebase": "workspace:*", "@genkit-ai/google-cloud": "workspace:*", - "@genkit-ai/googleai": "workspace:*", - "@genkit-ai/vertexai": "workspace:*", - "express": "^4.20.0" + "@genkit-ai/google-genai": "workspace:*", + "express": "^4.20.0", + "node-fetch": "3.3.2", + "wav": "^1.0.2" }, "devDependencies": { + "@types/wav": "^1.0.4", "typescript": "^5.6.2" } } diff --git a/js/testapps/basic-gemini/palm_tree.png b/js/testapps/basic-gemini/palm_tree.png new file mode 100644 index 0000000000..79388bf0f8 Binary files /dev/null and b/js/testapps/basic-gemini/palm_tree.png differ diff --git a/js/testapps/basic-gemini/photo.jpg b/js/testapps/basic-gemini/photo.jpg new file mode 100644 index 0000000000..fcf6dfa6ad Binary files /dev/null and b/js/testapps/basic-gemini/photo.jpg differ diff --git a/js/testapps/basic-gemini/src/index-vertexai.ts b/js/testapps/basic-gemini/src/index-vertexai.ts new file mode 100644 index 0000000000..dbf31da288 --- /dev/null +++ b/js/testapps/basic-gemini/src/index-vertexai.ts @@ -0,0 +1,296 @@ +/** + * 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 { vertexAI } from '@genkit-ai/google-genai'; +import * as fs from 'fs'; +import { genkit, z } from 'genkit'; +import wav from 'wav'; + +const ai = genkit({ + plugins: [ + // Make sure your Application Default Credentials are set + vertexAI({ experimental_debugTraces: true, location: 'global' }), + ], +}); + +// Basic Hi +ai.defineFlow('basic-hi', async () => { + const { text } = await ai.generate({ + model: vertexAI.model('gemini-2.5-flash'), + prompt: 'You are a helpful AI assistant named Walt, say hello', + }); + + return text; +}); + +// Multimodal input +ai.defineFlow('multimodal-input', async () => { + const photoBase64 = fs.readFileSync('photo.jpg', { encoding: 'base64' }); + + const { text } = await ai.generate({ + model: vertexAI.model('gemini-2.5-flash'), + prompt: [ + { text: 'describe this photo' }, + { + media: { + contentType: 'image/jpeg', + url: `data:image/jpeg;base64,${photoBase64}`, + }, + }, + ], + }); + + return text; +}); + +// YouTube videos +ai.defineFlow('youtube-videos', async (_, { sendChunk }) => { + const { text } = await ai.generate({ + model: vertexAI.model('gemini-2.5-flash'), + prompt: [ + { + text: 'transcribe this video', + }, + { + media: { + url: 'https://www.youtube.com/watch?v=3p1P5grjXIQ', + contentType: 'video/mp4', + }, + }, + ], + }); + + return text; +}); + +// streaming +ai.defineFlow('streaming', async (_, { sendChunk }) => { + const { stream } = ai.generateStream({ + model: vertexAI.model('gemini-2.5-flash'), + prompt: 'Write a poem about AI.', + }); + + let poem = ''; + for await (const chunk of stream) { + poem += chunk.text; + sendChunk(chunk.text); + } + + return poem; +}); + +// Search grounding +ai.defineFlow('search-grounding', async () => { + const { text, raw } = await ai.generate({ + model: vertexAI.model('gemini-2.5-flash'), + prompt: 'Who is Albert Einstein?', + config: { + tools: [{ googleSearch: {} }], + }, + }); + + return { + text, + groundingMetadata: (raw as any)?.candidates[0]?.groundingMetadata, + }; +}); + +const getWeather = ai.defineTool( + { + name: 'getWeather', + inputSchema: z.object({ + location: z + .string() + .describe( + 'Location for which to get the weather, ex: San-Francisco, CA' + ), + }), + description: 'used to get current weather for a location', + }, + async (input) => { + // pretend we call an actual API + return { + location: input.location, + temperature_celcius: 21.5, + conditions: 'cloudy', + }; + } +); + +const celsiusToFahrenheit = ai.defineTool( + { + name: 'celsiusToFahrenheit', + inputSchema: z.object({ + celsius: z.number().describe('Temperature in Celsius'), + }), + description: 'Converts Celsius to Fahrenheit', + }, + async ({ celsius }) => { + return (celsius * 9) / 5 + 32; + } +); + +// Tool calling with Gemini +ai.defineFlow( + { + name: 'toolCalling', + inputSchema: z.string().default('Paris, France'), + outputSchema: z.string(), + streamSchema: z.any(), + }, + async (location, { sendChunk }) => { + const { response, stream } = ai.generateStream({ + model: vertexAI.model('gemini-2.5-flash'), + config: { + temperature: 1, + }, + tools: [getWeather, celsiusToFahrenheit], + prompt: `What's the weather in ${location}? Convert the temperature to Fahrenheit.`, + }); + + for await (const chunk of stream) { + sendChunk(chunk); + } + + return (await response).text; + } +); + +const RpgCharacterSchema = z.object({ + name: z.string().describe('name of the character'), + backstory: z.string().describe("character's backstory, about a paragraph"), + weapons: z.array(z.string()), + class: z.enum(['RANGER', 'WIZZARD', 'TANK', 'HEALER', 'ENGINEER']), +}); + +// A simple example of structured output. +ai.defineFlow( + { + name: 'structured-output', + inputSchema: z.string().default('Glorb'), + outputSchema: RpgCharacterSchema, + }, + async (name, { sendChunk }) => { + const { response, stream } = ai.generateStream({ + model: vertexAI.model('gemini-2.5-flash'), + config: { + temperature: 2, // we want creativity + }, + output: { schema: RpgCharacterSchema }, + prompt: `Generate an RPC character called ${name}`, + }); + + for await (const chunk of stream) { + sendChunk(chunk.output); + } + + return (await response).output!; + } +); + +// Gemini reasoning example. +ai.defineFlow('reasoning', async (_, { sendChunk }) => { + const { message } = await ai.generate({ + prompt: 'what is heavier, one kilo of steel or one kilo of feathers', + model: vertexAI.model('gemini-2.5-pro'), + config: { + thinkingConfig: { + thinkingBudget: 1024, + includeThoughts: true, + }, + }, + onChunk: sendChunk, + }); + + return message; +}); + +// Image editing with Gemini. +ai.defineFlow('gemini-image-editing', async (_) => { + const plant = fs.readFileSync('palm_tree.png', { encoding: 'base64' }); + const room = fs.readFileSync('my_room.png', { encoding: 'base64' }); + + const { media } = await ai.generate({ + model: vertexAI.model('gemini-2.5-flash-image-preview'), + prompt: [ + { text: 'add the plant to my room' }, + { media: { url: `data:image/png;base64,${plant}` } }, + { media: { url: `data:image/png;base64,${room}` } }, + ], + config: { + responseModalities: ['TEXT', 'IMAGE'], + }, + }); + + return media; +}); + +// A simple example of image generation with Gemini. +ai.defineFlow('imagen-image-generation', async (_) => { + const { media } = await ai.generate({ + model: vertexAI.model('imagen-3.0-generate-002'), + prompt: `generate an image of a banana riding a bicycle`, + }); + + return media; +}); + +// Music generation with Lyria +ai.defineFlow('lyria-music-generation', async (_) => { + const { media } = await ai.generate({ + model: vertexAI.model('lyria-002'), + prompt: 'generate a relaxing song with piano and violin', + }); + + if (!media) { + throw new Error('no media returned'); + } + const audioBuffer = Buffer.from( + media.url.substring(media.url.indexOf(',') + 1), + 'base64' + ); + return { + media: 'data:audio/wav;base64,' + (await toWav(audioBuffer, 2, 48000)), + }; +}); + +async function toWav( + pcmData: Buffer, + channels = 1, + rate = 24000, + sampleWidth = 2 +): Promise { + return new Promise((resolve, reject) => { + // This code depends on `wav` npm library. + const writer = new wav.Writer({ + channels, + sampleRate: rate, + bitDepth: sampleWidth * 8, + }); + + let bufs = [] as any[]; + writer.on('error', reject); + writer.on('data', function (d) { + bufs.push(d); + }); + writer.on('end', function () { + resolve(Buffer.concat(bufs).toString('base64')); + }); + + writer.write(pcmData); + writer.end(); + }); +} diff --git a/js/testapps/basic-gemini/src/index.ts b/js/testapps/basic-gemini/src/index.ts index ea5ffe8898..86de8c6a15 100644 --- a/js/testapps/basic-gemini/src/index.ts +++ b/js/testapps/basic-gemini/src/index.ts @@ -14,45 +14,377 @@ * limitations under the License. */ -import { gemini15Flash, googleAI } from '@genkit-ai/googleai'; -import { vertexAI } from '@genkit-ai/vertexai'; -import { genkit, z } from 'genkit'; +import { googleAI } from '@genkit-ai/google-genai'; +import * as fs from 'fs'; +import { MediaPart, genkit, z } from 'genkit'; +import { Readable } from 'stream'; +import wav from 'wav'; const ai = genkit({ - plugins: [googleAI(), vertexAI()], + plugins: [ + // Provide the key via the GOOGLE_GENAI_API_KEY environment variable or arg { apiKey: 'yourkey'} + googleAI({ experimental_debugTraces: true }), + ], }); -const jokeSubjectGenerator = ai.defineTool( +ai.defineFlow('basic-hi', async () => { + const { text } = await ai.generate({ + model: googleAI.model('gemini-2.5-flash'), + prompt: 'You are a helpful AI assistant named Walt, say hello', + }); + + return text; +}); + +// Multimodal input +ai.defineFlow('multimodal-input', async () => { + const photoBase64 = fs.readFileSync('photo.jpg', { encoding: 'base64' }); + + const { text } = await ai.generate({ + model: googleAI.model('gemini-2.5-flash'), + prompt: [ + { text: 'describe this photo' }, + { + media: { + contentType: 'image/jpeg', + url: `data:image/jpeg;base64,${photoBase64}`, + }, + }, + ], + }); + + return text; +}); + +// YouTube videos +ai.defineFlow('youtube-videos', async (_, { sendChunk }) => { + const { text } = await ai.generate({ + model: googleAI.model('gemini-2.5-flash'), + prompt: [ + { + text: 'transcribe this video', + }, + { + media: { + url: 'https://www.youtube.com/watch?v=3p1P5grjXIQ', + contentType: 'video/mp4', + }, + }, + ], + }); + + return text; +}); + +// streaming +ai.defineFlow('streaming', async (_, { sendChunk }) => { + const { stream } = ai.generateStream({ + model: googleAI.model('gemini-2.5-flash'), + prompt: 'Write a poem about AI.', + }); + + let poem = ''; + for await (const chunk of stream) { + poem += chunk.text; + sendChunk(chunk.text); + } + + return poem; +}); + +// Search grounding +ai.defineFlow('search-grounding', async () => { + const { text, raw } = await ai.generate({ + model: googleAI.model('gemini-2.5-flash'), + prompt: 'Who is Albert Einstein?', + config: { + tools: [{ googleSearch: {} }], + }, + }); + + return { + text, + groundingMetadata: (raw as any)?.candidates[0]?.groundingMetadata, + }; +}); + +const getWeather = ai.defineTool( { - name: 'jokeSubjectGenerator', - description: 'Can be called to generate a subject for a joke', + name: 'getWeather', + inputSchema: z.object({ + location: z + .string() + .describe( + 'Location for which to get the weather, ex: San-Francisco, CA' + ), + }), + description: 'used to get current weather for a location', }, - async () => { - return 'banana'; + async (input) => { + // pretend we call an actual API + return { + location: input.location, + temperature_celcius: 21.5, + conditions: 'cloudy', + }; } ); -export const jokeFlow = ai.defineFlow( +const celsiusToFahrenheit = ai.defineTool( { - name: 'jokeFlow', - inputSchema: z.void(), - outputSchema: z.any(), + name: 'celsiusToFahrenheit', + inputSchema: z.object({ + celsius: z.number().describe('Temperature in Celsius'), + }), + description: 'Converts Celsius to Fahrenheit', }, - async () => { - const llmResponse = await ai.generate({ - stepName: 'joke-creator', - model: gemini15Flash, + async ({ celsius }) => { + return (celsius * 9) / 5 + 32; + } +); + +// Tool calling with Gemini +ai.defineFlow( + { + name: 'toolCalling', + inputSchema: z.string().default('Paris, France'), + outputSchema: z.string(), + streamSchema: z.any(), + }, + async (location, { sendChunk }) => { + const { response, stream } = ai.generateStream({ + model: googleAI.model('gemini-2.5-flash'), + config: { + temperature: 1, + }, + tools: [getWeather, celsiusToFahrenheit], + prompt: `What's the weather in ${location}? Convert the temperature to Fahrenheit.`, + }); + + for await (const chunk of stream) { + sendChunk(chunk); + } + + return (await response).text; + } +); + +const RpgCharacterSchema = z.object({ + name: z.string().describe('name of the character'), + backstory: z.string().describe("character's backstory, about a paragraph"), + weapons: z.array(z.string()), + class: z.enum(['RANGER', 'WIZZARD', 'TANK', 'HEALER', 'ENGINEER']), +}); + +// A simple example of structured output. +ai.defineFlow( + { + name: 'structured-output', + inputSchema: z.string().default('Glorb'), + outputSchema: RpgCharacterSchema, + }, + async (name, { sendChunk }) => { + const { response, stream } = ai.generateStream({ + model: googleAI.model('gemini-2.5-flash'), config: { - temperature: 2, - // if desired, model versions can be explicitly set - version: 'gemini-1.5-flash-002', + temperature: 2, // we want creativity + }, + output: { schema: RpgCharacterSchema }, + prompt: `Generate an RPC character called ${name}`, + }); + + for await (const chunk of stream) { + sendChunk(chunk.output); + } + + return (await response).output!; + } +); + +// Gemini reasoning example. +ai.defineFlow('reasoning', async (_, { sendChunk }) => { + const { message } = await ai.generate({ + prompt: 'what is heavier, one kilo of steel or one kilo of feathers', + model: googleAI.model('gemini-2.5-pro'), + config: { + thinkingConfig: { + thinkingBudget: 1024, + includeThoughts: true, }, - output: { - schema: z.object({ jokeSubject: z.string() }), + }, + onChunk: sendChunk, + }); + + return message; +}); + +// Image editing with Gemini. +ai.defineFlow('gemini-image-editing', async (_) => { + const plant = fs.readFileSync('palm_tree.png', { encoding: 'base64' }); + const room = fs.readFileSync('my_room.png', { encoding: 'base64' }); + + const { media } = await ai.generate({ + model: googleAI.model('gemini-2.5-flash-image-preview'), + prompt: [ + { text: 'add the plant to my room' }, + { media: { url: `data:image/png;base64,${plant}` } }, + { media: { url: `data:image/png;base64,${room}` } }, + ], + config: { + responseModalities: ['TEXT', 'IMAGE'], + }, + }); + + return media; +}); + +// A simple example of image generation with Gemini. +ai.defineFlow('imagen-image-generation', async (_) => { + const { media } = await ai.generate({ + model: googleAI.model('imagen-3.0-generate-002'), + prompt: `generate an image of a banana riding a bicycle`, + }); + + return media; +}); + +// TTS sample +ai.defineFlow( + { + name: 'tts', + inputSchema: z + .string() + .default( + 'Gemini is amazing. Can say things like: glorg, blub-blub, and ayeeeeee!!!' + ), + outputSchema: z.object({ media: z.string() }), + }, + async (prompt) => { + const { media } = await ai.generate({ + model: googleAI.model('gemini-2.5-flash-preview-tts'), + config: { + responseModalities: ['AUDIO'], + // For all available options see https://ai.google.dev/gemini-api/docs/speech-generation#javascript + speechConfig: { + voiceConfig: { + prebuiltVoiceConfig: { voiceName: 'Algenib' }, + }, + }, }, - tools: [jokeSubjectGenerator], - prompt: `come up with a subject to joke about (using the function provided)`, + prompt, }); - return llmResponse.output; + if (!media) { + throw new Error('no media returned'); + } + const audioBuffer = Buffer.from( + media.url.substring(media.url.indexOf(',') + 1), + 'base64' + ); + return { + media: 'data:audio/wav;base64,' + (await toWav(audioBuffer)), + }; } ); + +async function toWav( + pcmData: Buffer, + channels = 1, + rate = 24000, + sampleWidth = 2 +): Promise { + return new Promise((resolve, reject) => { + // This code depends on `wav` npm library. + const writer = new wav.Writer({ + channels, + sampleRate: rate, + bitDepth: sampleWidth * 8, + }); + + let bufs = [] as any[]; + writer.on('error', reject); + writer.on('data', function (d) { + bufs.push(d); + }); + writer.on('end', function () { + resolve(Buffer.concat(bufs).toString('base64')); + }); + + writer.write(pcmData); + writer.end(); + }); +} + +// An example of using Ver 2 model to make a static photo move. +ai.defineFlow('photo-move-veo', async (_, { sendChunk }) => { + const startingImage = fs.readFileSync('photo.jpg', { encoding: 'base64' }); + + let { operation } = await ai.generate({ + model: googleAI.model('veo-2.0-generate-001'), + prompt: [ + { + text: 'make the subject in the photo move', + }, + { + media: { + contentType: 'image/jpeg', + url: `data:image/jpeg;base64,${startingImage}`, + }, + }, + ], + config: { + durationSeconds: 5, + aspectRatio: '9:16', + personGeneration: 'allow_adult', + }, + }); + + if (!operation) { + throw new Error('Expected the model to return an operation'); + } + + while (!operation.done) { + sendChunk('check status of operation ' + operation.id); + operation = await ai.checkOperation(operation); + await new Promise((resolve) => setTimeout(resolve, 5000)); + } + + if (operation.error) { + sendChunk('Error: ' + operation.error.message); + throw new Error('failed to generate video: ' + operation.error.message); + } + + // operation done, download generated video to disk + const video = operation.output?.message?.content.find((p) => !!p.media); + if (!video) { + throw new Error('Failed to find the generated video'); + } + sendChunk('Writing results to photo.mp4'); + await downloadVideo(video, 'photo.mp4'); + sendChunk('Done!'); + + return operation; +}); + +function getApiKeyFromEnvVar(): string | undefined { + return ( + process.env.GEMINI_API_KEY || + process.env.GOOGLE_API_KEY || + process.env.GOOGLE_GENAI_API_KEY + ); +} + +async function downloadVideo(video: MediaPart, path: string) { + const fetch = (await import('node-fetch')).default; + const videoDownloadResponse = await fetch( + `${video.media!.url}&key=${getApiKeyFromEnvVar()}` + ); + if ( + !videoDownloadResponse || + videoDownloadResponse.status !== 200 || + !videoDownloadResponse.body + ) { + throw new Error('Failed to fetch video'); + } + + Readable.from(videoDownloadResponse.body).pipe(fs.createWriteStream(path)); +} diff --git a/samples/js-gemini/src/index.ts b/samples/js-gemini/src/index.ts index 6117685e43..86de8c6a15 100644 --- a/samples/js-gemini/src/index.ts +++ b/samples/js-gemini/src/index.ts @@ -16,7 +16,7 @@ import { googleAI } from '@genkit-ai/google-genai'; import * as fs from 'fs'; -import { genkit, MediaPart, z } from 'genkit'; +import { MediaPart, genkit, z } from 'genkit'; import { Readable } from 'stream'; import wav from 'wav'; diff --git a/samples/package.json b/samples/package.json index a1ccfcee47..a5e8b2f322 100644 --- a/samples/package.json +++ b/samples/package.json @@ -10,6 +10,7 @@ "build:js-menu": "cd js-menu && npm install && npm run build", "build:js-prompts": "cd js-prompts && npm install && npm run build", "build:js-schoolAgent": "cd js-schoolAgent && npm install && npm run build", + "build:js-gemini": "cd js-gemini && npm install && npm run build", "build:all-samples": "concurrently npm:build:js-*" }, "pre-commit": [