From 1cbd3fbe04de34969d8968e9aeacb3f1bb6c455f Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Tue, 26 Aug 2025 09:14:41 +0100 Subject: [PATCH 1/7] refactor(plugins/googleai): migrate to v2 API --- js/plugins/googleai/src/embedder.ts | 29 ++- js/plugins/googleai/src/gemini.ts | 103 ++++------ js/plugins/googleai/src/imagen.ts | 12 +- js/plugins/googleai/src/index.ts | 244 ++++++++++------------- js/plugins/googleai/src/veo.ts | 12 +- js/plugins/googleai/tests/gemini_test.ts | 75 +++---- 6 files changed, 194 insertions(+), 281 deletions(-) diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index 531ac8a0e5..e1b219f7d6 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -23,9 +23,9 @@ import { z, type EmbedderAction, type EmbedderReference, - type Genkit, } from 'genkit'; import { embedderRef } from 'genkit/embedder'; +import { embedder } from 'genkit/plugin'; import { getApiKeyFromEnvVar } from './common.js'; import type { PluginOptions } from './index.js'; @@ -97,13 +97,12 @@ export const geminiEmbedding001 = embedderRef({ }); export const SUPPORTED_MODELS = { - 'embedding-001': textEmbeddingGecko001, - 'text-embedding-004': textEmbedding004, - 'gemini-embedding-001': geminiEmbedding001, + 'googleai/embedding-001': textEmbeddingGecko001, + 'googleai/text-embedding-004': textEmbedding004, + 'googleai/gemini-embedding-001': geminiEmbedding001, }; export function defineGoogleAIEmbedder( - ai: Genkit, name: string, pluginOptions: PluginOptions ): EmbedderAction { @@ -117,7 +116,7 @@ export function defineGoogleAIEmbedder( 'For more details see https://genkit.dev/docs/plugins/google-genai' ); } - const embedder: EmbedderReference = + const embedderReference: EmbedderReference = SUPPORTED_MODELS[name] ?? embedderRef({ name: name, @@ -130,16 +129,16 @@ export function defineGoogleAIEmbedder( }, }, }); - const apiModelName = embedder.name.startsWith('googleai/') - ? embedder.name.substring('googleai/'.length) - : embedder.name; - return ai.defineEmbedder( + const apiModelName = embedderReference.name.startsWith('googleai/') + ? embedderReference.name.substring('googleai/'.length) + : embedderReference.name; + return embedder( { - name: embedder.name, + name: embedderReference.name, configSchema: GeminiEmbeddingConfigSchema, - info: embedder.info!, + info: embedderReference.info!, }, - async (input, options) => { + async ({ input, options }) => { if (pluginOptions.apiKey === false && !options?.apiKey) { throw new GenkitError({ status: 'INVALID_ARGUMENT', @@ -152,8 +151,8 @@ export function defineGoogleAIEmbedder( ).getGenerativeModel({ model: options?.version || - embedder.config?.version || - embedder.version || + embedderReference.config?.version || + embedderReference.version || apiModelName, }); const embeddings = await Promise.all( diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index 38d8285856..33a0550f35 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -38,7 +38,7 @@ import { type ToolConfig, type UsageMetadata, } from '@google/generative-ai'; -import { GenkitError, z, type Genkit, type JSONSchema } from 'genkit'; +import { GenkitError, z, type JSONSchema } from 'genkit'; import { GenerationCommonConfigDescriptions, GenerationCommonConfigSchema, @@ -57,8 +57,8 @@ import { type ToolResponsePart, } from 'genkit/model'; import { downloadRequestMedia } from 'genkit/model/middleware'; -import { runInNewSpan } from 'genkit/tracing'; -import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common'; +import { model } from 'genkit/plugin'; +import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common.js'; import { handleCacheIfNeeded } from './context-caching'; import { extractCacheConfig } from './context-caching/utils'; @@ -644,26 +644,26 @@ export const gemma3ne4bit = modelRef({ }); export const SUPPORTED_GEMINI_MODELS = { - 'gemini-1.5-pro': gemini15Pro, - 'gemini-1.5-flash': gemini15Flash, - 'gemini-1.5-flash-8b': gemini15Flash8b, - 'gemini-2.0-pro-exp-02-05': gemini20ProExp0205, - 'gemini-2.0-flash': gemini20Flash, - 'gemini-2.0-flash-lite': gemini20FlashLite, - 'gemini-2.0-flash-exp': gemini20FlashExp, - 'gemini-2.5-pro-exp-03-25': gemini25ProExp0325, - 'gemini-2.5-pro-preview-03-25': gemini25ProPreview0325, - 'gemini-2.5-pro-preview-tts': gemini25ProPreviewTts, - 'gemini-2.5-flash-preview-04-17': gemini25FlashPreview0417, - 'gemini-2.5-flash-preview-tts': gemini25FlashPreviewTts, - 'gemini-2.5-flash': gemini25Flash, - 'gemini-2.5-flash-lite': gemini25FlashLite, - 'gemini-2.5-pro': gemini25Pro, - 'gemma-3-12b-it': gemma312bit, - 'gemma-3-1b-it': gemma31bit, - 'gemma-3-27b-it': gemma327bit, - 'gemma-3-4b-it': gemma34bit, - 'gemma-3n-e4b-it': gemma3ne4bit, + 'googleai/gemini-1.5-pro': gemini15Pro, + 'googleai/gemini-1.5-flash': gemini15Flash, + 'googleai/gemini-1.5-flash-8b': gemini15Flash8b, + 'googleai/gemini-2.0-pro-exp-02-05': gemini20ProExp0205, + 'googleai/gemini-2.0-flash': gemini20Flash, + 'googleai/gemini-2.0-flash-lite': gemini20FlashLite, + 'googleai/gemini-2.0-flash-exp': gemini20FlashExp, + 'googleai/gemini-2.5-pro-exp-03-25': gemini25ProExp0325, + 'googleai/gemini-2.5-pro-preview-03-25': gemini25ProPreview0325, + 'googleai/gemini-2.5-pro-preview-tts': gemini25ProPreviewTts, + 'googleai/gemini-2.5-flash-preview-04-17': gemini25FlashPreview0417, + 'googleai/gemini-2.5-flash-preview-tts': gemini25FlashPreviewTts, + 'googleai/gemini-2.5-flash': gemini25Flash, + 'googleai/gemini-2.5-flash-lite': gemini25FlashLite, + 'googleai/gemini-2.5-pro': gemini25Pro, + 'googleai/gemma-3-12b-it': gemma312bit, + 'googleai/gemma-3-1b-it': gemma31bit, + 'googleai/gemma-3-27b-it': gemma327bit, + 'googleai/gemma-3-4b-it': gemma34bit, + 'googleai/gemma-3n-e4b-it': gemma3ne4bit, }; export const GENERIC_GEMINI_MODEL = modelRef({ @@ -705,7 +705,7 @@ export type GeminiVersionString = * ```js * await ai.generate({ * prompt: 'hi', - * model: gemini('gemini-1.5-flash') + * model: gemini('googleai/gemini-1.5-flash') * }); * ``` */ @@ -1118,7 +1118,6 @@ export function cleanSchema(schema: JSONSchema): JSONSchema { * Defines a new GoogleAI model. */ export function defineGoogleAIModel({ - ai, name, apiKey: apiKeyOption, apiVersion, @@ -1127,7 +1126,6 @@ export function defineGoogleAIModel({ defaultConfig, debugTraces, }: { - ai: Genkit; name: string; apiKey?: string | false; apiVersion?: string; @@ -1154,10 +1152,10 @@ export function defineGoogleAIModel({ ? name.substring('googleai/'.length) : name; - const model: ModelReference = + const modelReference: ModelReference = SUPPORTED_GEMINI_MODELS[apiModelName] ?? modelRef({ - name: `googleai/${apiModelName}`, + name: name, // Keep the full name for the model reference info: { label: `Google AI - ${apiModelName}`, supports: { @@ -1173,7 +1171,7 @@ export function defineGoogleAIModel({ }); const middleware: ModelMiddleware[] = []; - if (model.info?.supports?.media) { + if (modelReference.info?.supports?.media) { // the gemini api doesn't support downloading media from http(s) middleware.push( downloadRequestMedia({ @@ -1199,12 +1197,11 @@ export function defineGoogleAIModel({ ); } - return ai.defineModel( + return model( { - apiVersion: 'v2', - name: model.name, - ...model.info, - configSchema: model.configSchema, + name: modelReference.name, + ...modelReference.info, + configSchema: modelReference.configSchema, use: middleware, }, async (request, { streamingRequested, sendChunk, abortSignal }) => { @@ -1228,7 +1225,7 @@ export function defineGoogleAIModel({ // systemInstructions to be provided as a separate input. The first // message detected with role=system will be used for systemInstructions. let systemInstruction: GeminiMessage | undefined = undefined; - if (model.info?.supports?.systemRole) { + if (modelReference.info?.supports?.systemRole) { const systemMessage = messages.find((m) => m.role === 'system'); if (systemMessage) { messages.splice(messages.indexOf(systemMessage), 1); @@ -1306,7 +1303,10 @@ export function defineGoogleAIModel({ generationConfig.responseSchema = cleanSchema(request.output.schema); } - const msg = toGeminiMessage(messages[messages.length - 1], model); + const msg = toGeminiMessage( + messages[messages.length - 1], + modelReference + ); const fromJSONModeScopedGeminiCandidate = ( candidate: GeminiCandidate @@ -1321,11 +1321,11 @@ export function defineGoogleAIModel({ toolConfig, history: messages .slice(0, -1) - .map((message) => toGeminiMessage(message, model)), + .map((message) => toGeminiMessage(message, modelReference)), safetySettings: safetySettingsFromConfig, } as StartChatParams; const modelVersion = (versionFromConfig || - model.version || + modelReference.version || apiModelName) as string; const cacheConfigDetails = extractCacheConfig(request); @@ -1426,31 +1426,8 @@ export function defineGoogleAIModel({ }; }; - // If debugTraces is enable, we wrap the actual model call with a span, add raw - // API params as for input. - return debugTraces - ? await runInNewSpan( - ai.registry, - { - metadata: { - name: streamingRequested ? 'sendMessageStream' : 'sendMessage', - }, - }, - async (metadata) => { - metadata.input = { - sdk: '@google/generative-ai', - cache: cache, - model: genModel.model, - chatOptions: updatedChatRequest, - parts: msg.parts, - options, - }; - const response = await callGemini(); - metadata.output = response.custom; - return response; - } - ) - : await callGemini(); + // TODO v2: no ai.registry available here; run without the debug span wrapper. + return await callGemini(); } ); } diff --git a/js/plugins/googleai/src/imagen.ts b/js/plugins/googleai/src/imagen.ts index d7d6c39709..335c4649aa 100644 --- a/js/plugins/googleai/src/imagen.ts +++ b/js/plugins/googleai/src/imagen.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { GenkitError, MessageData, z, type Genkit } from 'genkit'; +import { GenkitError, MessageData, z } from 'genkit'; import { getBasicUsageStats, modelRef, @@ -23,6 +23,7 @@ import { type ModelInfo, type ModelReference, } from 'genkit/model'; +import { model } from 'genkit/plugin'; import { getApiKeyFromEnvVar } from './common.js'; import { predictModel } from './predict.js'; @@ -109,7 +110,6 @@ export const GENERIC_IMAGEN_INFO = { } as ModelInfo; export function defineImagenModel( - ai: Genkit, name: string, apiKey?: string | false ): ModelAction { @@ -125,7 +125,7 @@ export function defineImagenModel( } } const modelName = `googleai/${name}`; - const model: ModelReference = modelRef({ + const modelReference: ModelReference = modelRef({ name: modelName, info: { ...GENERIC_IMAGEN_INFO, @@ -134,10 +134,10 @@ export function defineImagenModel( configSchema: ImagenConfigSchema, }); - return ai.defineModel( + return model( { name: modelName, - ...model.info, + ...modelReference.info, configSchema: ImagenConfigSchema, }, async (request) => { @@ -153,7 +153,7 @@ export function defineImagenModel( ImagenInstance, ImagenPrediction, ImagenParameters - >(model.version || name, apiKey as string, 'predict'); + >(modelReference.version || name, apiKey as string, 'predict'); const response = await predictClient([instance], toParameters(request)); if (!response.predictions || response.predictions.length == 0) { diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts index 52cb48a880..514b1e83f9 100644 --- a/js/plugins/googleai/src/index.ts +++ b/js/plugins/googleai/src/index.ts @@ -20,13 +20,12 @@ import { modelActionMetadata, type ActionMetadata, type EmbedderReference, - type Genkit, type ModelReference, type z, } from 'genkit'; import { logger } from 'genkit/logging'; import { modelRef } from 'genkit/model'; -import { genkitPlugin, type GenkitPlugin } from 'genkit/plugin'; +import { genkitPluginV2, type GenkitPluginV2 } from 'genkit/plugin'; import type { ActionType } from 'genkit/registry'; import { getApiKeyFromEnvVar } from './common.js'; import { @@ -112,119 +111,6 @@ export interface PluginOptions { experimental_debugTraces?: boolean; } -async function initializer(ai: Genkit, options?: PluginOptions) { - let apiVersions = ['v1']; - - if (options?.apiVersion) { - if (Array.isArray(options?.apiVersion)) { - apiVersions = options?.apiVersion; - } else { - apiVersions = [options?.apiVersion]; - } - } - - if (apiVersions.includes('v1beta')) { - Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => - defineGoogleAIModel({ - ai, - name, - apiKey: options?.apiKey, - apiVersion: 'v1beta', - baseUrl: options?.baseUrl, - debugTraces: options?.experimental_debugTraces, - }) - ); - } - if (apiVersions.includes('v1')) { - Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => - defineGoogleAIModel({ - ai, - name, - apiKey: options?.apiKey, - apiVersion: undefined, - baseUrl: options?.baseUrl, - debugTraces: options?.experimental_debugTraces, - }) - ); - Object.keys(EMBEDDER_MODELS).forEach((name) => - defineGoogleAIEmbedder(ai, name, { apiKey: options?.apiKey }) - ); - } - - if (options?.models) { - for (const modelOrRef of options?.models) { - const modelName = - typeof modelOrRef === 'string' - ? modelOrRef - : // strip out the `googleai/` prefix - modelOrRef.name.split('/')[1]; - const modelRef = - typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef; - defineGoogleAIModel({ - ai, - name: modelName, - apiKey: options?.apiKey, - baseUrl: options?.baseUrl, - info: { - ...modelRef.info, - label: `Google AI - ${modelName}`, - }, - debugTraces: options?.experimental_debugTraces, - }); - } - } -} - -async function resolver( - ai: Genkit, - actionType: ActionType, - actionName: string, - options?: PluginOptions -) { - if (actionType === 'embedder') { - resolveEmbedder(ai, actionName, options); - } else if (actionName.startsWith('veo')) { - // we do it this way because the request may come in for - // action type 'model' and action name 'veo-...'. That case should - // be a noop. It's just the order or model lookup. - if (actionType === 'background-model') { - defineVeoModel(ai, actionName, options?.apiKey); - } - } else if (actionType === 'model') { - resolveModel(ai, actionName, options); - } -} - -function resolveModel(ai: Genkit, actionName: string, options?: PluginOptions) { - if (actionName.startsWith('imagen')) { - defineImagenModel(ai, actionName, options?.apiKey); - return; - } - - const modelRef = gemini(actionName); - defineGoogleAIModel({ - ai, - name: modelRef.name, - apiKey: options?.apiKey, - baseUrl: options?.baseUrl, - info: { - ...modelRef.info, - label: `Google AI - ${actionName}`, - }, - debugTraces: options?.experimental_debugTraces, - }); -} - -function resolveEmbedder( - ai: Genkit, - actionName: string, - options?: PluginOptions -) { - defineGoogleAIEmbedder(ai, `googleai/${actionName}`, { - apiKey: options?.apiKey, - }); -} - async function listActions(options?: PluginOptions): Promise { const apiKey = options?.apiKey || getApiKeyFromEnvVar(); if (!apiKey) { @@ -256,7 +142,7 @@ async function listActions(options?: PluginOptions): Promise { const name = m.name.split('/').at(-1)!; return modelActionMetadata({ - name: `googleai/${name}`, + name: name, info: { ...GENERIC_IMAGEN_INFO }, configSchema: ImagenConfigSchema, }); @@ -274,7 +160,7 @@ async function listActions(options?: PluginOptions): Promise { const name = m.name.split('/').at(-1)!; return modelActionMetadata({ - name: `googleai/${name}`, + name: name, // Return unprefixed name for v2 info: { ...GENERIC_VEO_INFO }, configSchema: VeoConfigSchema, background: true, @@ -286,14 +172,13 @@ async function listActions(options?: PluginOptions): Promise { // Filter out deprecated .filter((m) => !m.description || !m.description.includes('deprecated')) .map((m) => { - const ref = gemini( - m.name.startsWith('models/') - ? m.name.substring('models/'.length) - : m.name - ); + const bare = m.name.startsWith('models/') + ? m.name.substring('models/'.length) + : m.name; + const ref = gemini(bare); return modelActionMetadata({ - name: ref.name, + name: bare, // Return unprefixed name for v2 info: ref.info, configSchema: GeminiConfigSchema, }); @@ -304,18 +189,16 @@ async function listActions(options?: PluginOptions): Promise { // Filter out deprecated .filter((m) => !m.description || !m.description.includes('deprecated')) .map((m) => { - const name = - 'googleai/' + - (m.name.startsWith('models/') - ? m.name.substring('models/'.length) - : m.name); + const bare = m.name.startsWith('models/') + ? m.name.substring('models/'.length) + : m.name; return embedderActionMetadata({ - name, + name: bare, // Return unprefixed name for v2 configSchema: GeminiEmbeddingConfigSchema, info: { dimensions: 768, - label: `Google Gen AI - ${name}`, + label: `Google Gen AI - ${bare}`, supports: { input: ['text'], }, @@ -328,23 +211,104 @@ async function listActions(options?: PluginOptions): Promise { /** * Google Gemini Developer API plugin. */ -export function googleAIPlugin(options?: PluginOptions): GenkitPlugin { +export function googleAIPlugin(options?: PluginOptions): 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', + async init() { + // Eagerly return actions to register. + const actions: any[] = []; + const apiVersions = Array.isArray(options?.apiVersion) + ? options!.apiVersion + : options?.apiVersion + ? [options.apiVersion] + : ['v1']; + + if (apiVersions.includes('v1beta')) { + Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { + actions.push( + defineGoogleAIModel({ + name, + apiKey: options?.apiKey, + apiVersion: 'v1beta', + baseUrl: options?.baseUrl, + debugTraces: options?.experimental_debugTraces, + }) + ); + }); + } + + if (apiVersions.includes('v1')) { + Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { + actions.push( + defineGoogleAIModel({ + name, + apiKey: options?.apiKey, + baseUrl: options?.baseUrl, + debugTraces: options?.experimental_debugTraces, + }) + ); + }); + Object.keys(EMBEDDER_MODELS).forEach((name) => { + actions.push( + defineGoogleAIEmbedder(name, { apiKey: options?.apiKey }) + ); + }); + } + + if (options?.models) { + for (const modelOrRef of options.models) { + const modelName = + typeof modelOrRef === 'string' + ? modelOrRef + : modelOrRef.name.split('/').pop()!; + const ref = + typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef; + actions.push( + defineGoogleAIModel({ + name: modelName, + apiKey: options?.apiKey, + baseUrl: options?.baseUrl, + info: { ...ref.info, label: `Google AI - ${modelName}` }, + debugTraces: options?.experimental_debugTraces, + }) + ); + } + } + return actions; + }, + async resolve(actionType: ActionType, actionName: string) { + if (actionType === 'embedder') { + return defineGoogleAIEmbedder(actionName, { apiKey: options?.apiKey }); + } else if (actionName.startsWith('veo')) { + if (actionType === 'background-model') { + return defineVeoModel(actionName, options?.apiKey); + } + } else if (actionType === 'model') { + if (actionName.startsWith('imagen')) { + return defineImagenModel(actionName, options?.apiKey); + } + // Fallback dynamic Gemini model + return defineGoogleAIModel({ + name: actionName, + apiKey: options?.apiKey, + baseUrl: options?.baseUrl, + debugTraces: options?.experimental_debugTraces, + }); + } + return undefined; + }, + async list() { if (listActionsCache) return listActionsCache; + // Must return UNNAMESPACED names in v2. Genkit will prefix them. listActionsCache = await listActions(options); return listActionsCache; - } - ); + }, + }); } export type GoogleAIPlugin = { - (params?: PluginOptions): GenkitPlugin; + (params?: PluginOptions): GenkitPluginV2; model( name: keyof typeof SUPPORTED_GEMINI_MODELS | (`gemini-${string}` & {}), config?: z.infer diff --git a/js/plugins/googleai/src/veo.ts b/js/plugins/googleai/src/veo.ts index b35c87a377..be1dba1bc7 100644 --- a/js/plugins/googleai/src/veo.ts +++ b/js/plugins/googleai/src/veo.ts @@ -14,13 +14,7 @@ * limitations under the License. */ -import { - GenerateResponseData, - GenkitError, - Operation, - z, - type Genkit, -} from 'genkit'; +import { GenerateResponseData, GenkitError, Operation, z } from 'genkit'; import { BackgroundModelAction, modelRef, @@ -28,6 +22,7 @@ import { type ModelInfo, type ModelReference, } from 'genkit/model'; +import { backgroundModel } from 'genkit/plugin'; import { getApiKeyFromEnvVar } from './common.js'; import { Operation as ApiOperation, checkOp, predictModel } from './predict.js'; @@ -127,7 +122,6 @@ export const GENERIC_VEO_INFO = { } as ModelInfo; export function defineVeoModel( - ai: Genkit, name: string, apiKey?: string | false ): BackgroundModelAction { @@ -152,7 +146,7 @@ export function defineVeoModel( configSchema: VeoConfigSchema, }); - return ai.defineBackgroundModel({ + return backgroundModel({ name: modelName, ...model.info, configSchema: VeoConfigSchema, diff --git a/js/plugins/googleai/tests/gemini_test.ts b/js/plugins/googleai/tests/gemini_test.ts index 658adc7d90..1a3cefa098 100644 --- a/js/plugins/googleai/tests/gemini_test.ts +++ b/js/plugins/googleai/tests/gemini_test.ts @@ -574,20 +574,13 @@ describe('plugin', () => { }); it('references dynamic models', async () => { - const ai = genkit({ - plugins: [googleAI({})], - }); const giraffeRef = gemini('gemini-4.5-giraffe'); assert.strictEqual(giraffeRef.name, 'googleai/gemini-4.5-giraffe'); - const giraffe = await ai.registry.lookupAction( - `/model/${giraffeRef.name}` - ); - assert.ok(giraffe); - assert.strictEqual(giraffe.__action.name, 'googleai/gemini-4.5-giraffe'); + assertEqualModelInfo( - giraffe.__action.metadata?.model, + giraffeRef.info!, 'Google AI - gemini-4.5-giraffe', - GENERIC_GEMINI_MODEL.info! // <---- generic model fallback + GENERIC_GEMINI_MODEL.info! ); }); @@ -603,54 +596,40 @@ describe('plugin', () => { const pro002Ref = gemini('gemini-1.5-pro-002'); assert.strictEqual(pro002Ref.name, 'googleai/gemini-1.5-pro-002'); - assertEqualModelInfo( - pro002Ref.info!, - 'Google AI - gemini-1.5-pro-002', - gemini15Pro.info! + // Test the actual model info structure that's returned + assert.strictEqual( + pro002Ref.info!.label, + 'Google AI - gemini-1.5-pro-002' ); - const pro002 = await ai.registry.lookupAction(`/model/${pro002Ref.name}`); - assert.ok(pro002); - assert.strictEqual(pro002.__action.name, 'googleai/gemini-1.5-pro-002'); - assertEqualModelInfo( - pro002.__action.metadata?.model, - 'Google AI - gemini-1.5-pro-002', - gemini15Pro.info! + assert.deepStrictEqual( + pro002Ref.info!.supports, + gemini15Pro.info!.supports ); assert.strictEqual(flash002Ref.name, 'googleai/gemini-1.5-flash-002'); - assertEqualModelInfo( - flash002Ref.info!, - 'Google AI - gemini-1.5-flash-002', - gemini15Flash.info! - ); - const flash002 = await ai.registry.lookupAction( - `/model/${flash002Ref.name}` - ); - assert.ok(flash002); assert.strictEqual( - flash002.__action.name, - 'googleai/gemini-1.5-flash-002' - ); - assertEqualModelInfo( - flash002.__action.metadata?.model, - 'Google AI - gemini-1.5-flash-002', - gemini15Flash.info! + flash002Ref.info!.label, + 'Google AI - gemini-1.5-flash-002' ); + assert.deepStrictEqual(flash002Ref.info!.supports, { + multiturn: true, + media: true, + tools: true, + toolChoice: true, + systemRole: true, + constrained: 'no-tools', + }); + const bananaRef = gemini('gemini-4.0-banana'); assert.strictEqual(bananaRef.name, 'googleai/gemini-4.0-banana'); - assertEqualModelInfo( - bananaRef.info!, - 'Google AI - gemini-4.0-banana', - GENERIC_GEMINI_MODEL.info! // <---- generic model fallback + assert.strictEqual( + bananaRef.info!.label, + 'Google AI - gemini-4.0-banana' ); - const banana = await ai.registry.lookupAction(`/model/${bananaRef.name}`); - assert.ok(banana); - assert.strictEqual(banana.__action.name, 'googleai/gemini-4.0-banana'); - assertEqualModelInfo( - banana.__action.metadata?.model, - 'Google AI - gemini-4.0-banana', - GENERIC_GEMINI_MODEL.info! // <---- generic model fallback + assert.deepStrictEqual( + bananaRef.info!.supports, + GENERIC_GEMINI_MODEL.info!.supports ); }); }); From 2a17d1362b2d0ce166b70a0faeb3a112bf850e3f Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Tue, 26 Aug 2025 13:17:13 +0100 Subject: [PATCH 2/7] feat(plugins/googlai): restore debugTraces in v2 plugin migration --- js/plugins/googleai/src/gemini.ts | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index 33a0550f35..500542fa65 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -58,6 +58,7 @@ import { } from 'genkit/model'; import { downloadRequestMedia } from 'genkit/model/middleware'; import { model } from 'genkit/plugin'; +import { runInNewSpan } from 'genkit/tracing'; import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common.js'; import { handleCacheIfNeeded } from './context-caching'; import { extractCacheConfig } from './context-caching/utils'; @@ -1426,8 +1427,30 @@ export function defineGoogleAIModel({ }; }; - // TODO v2: no ai.registry available here; run without the debug span wrapper. - return await callGemini(); + // If debugTraces is enabled, we wrap the actual model call with a span, add raw + // API params as input. + return debugTraces + ? await runInNewSpan( + { + metadata: { + name: streamingRequested ? 'sendMessageStream' : 'sendMessage', + }, + }, + async (metadata) => { + metadata.input = { + sdk: '@google/generative-ai', + cache: cache, + model: genModel.model, + chatOptions: updatedChatRequest, + parts: msg.parts, + options, + }; + const response = await callGemini(); + metadata.output = response.custom; + return response; + } + ) + : await callGemini(); } ); } From db7a169d640f3db44507c39ba9eb76ccda259aa0 Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Tue, 26 Aug 2025 13:28:41 +0100 Subject: [PATCH 3/7] chore(plugins/googleai): better variable names --- js/plugins/googleai/src/index.ts | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts index 514b1e83f9..ca054ceba8 100644 --- a/js/plugins/googleai/src/index.ts +++ b/js/plugins/googleai/src/index.ts @@ -139,10 +139,10 @@ async function listActions(options?: PluginOptions): Promise { // Filter out deprecated .filter((m) => !m.description || !m.description.includes('deprecated')) .map((m) => { - const name = m.name.split('/').at(-1)!; + const imagenModelName = m.name.split('/').at(-1)!; return modelActionMetadata({ - name: name, + name: imagenModelName, info: { ...GENERIC_IMAGEN_INFO }, configSchema: ImagenConfigSchema, }); @@ -157,10 +157,10 @@ async function listActions(options?: PluginOptions): Promise { // Filter out deprecated .filter((m) => !m.description || !m.description.includes('deprecated')) .map((m) => { - const name = m.name.split('/').at(-1)!; + const veoModelName = m.name.split('/').at(-1)!; return modelActionMetadata({ - name: name, // Return unprefixed name for v2 + name: veoModelName, info: { ...GENERIC_VEO_INFO }, configSchema: VeoConfigSchema, background: true, @@ -172,13 +172,13 @@ async function listActions(options?: PluginOptions): Promise { // Filter out deprecated .filter((m) => !m.description || !m.description.includes('deprecated')) .map((m) => { - const bare = m.name.startsWith('models/') + const modelName = m.name.startsWith('models/') ? m.name.substring('models/'.length) : m.name; - const ref = gemini(bare); + const ref = gemini(modelName); return modelActionMetadata({ - name: bare, // Return unprefixed name for v2 + name: modelName, info: ref.info, configSchema: GeminiConfigSchema, }); @@ -189,16 +189,16 @@ async function listActions(options?: PluginOptions): Promise { // Filter out deprecated .filter((m) => !m.description || !m.description.includes('deprecated')) .map((m) => { - const bare = m.name.startsWith('models/') + const embedderName = m.name.startsWith('models/') ? m.name.substring('models/'.length) : m.name; return embedderActionMetadata({ - name: bare, // Return unprefixed name for v2 + name: embedderName, configSchema: GeminiEmbeddingConfigSchema, info: { dimensions: 768, - label: `Google Gen AI - ${bare}`, + label: `Google Gen AI - ${embedderName}`, supports: { input: ['text'], }, @@ -216,7 +216,6 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { return genkitPluginV2({ name: 'googleai', async init() { - // Eagerly return actions to register. const actions: any[] = []; const apiVersions = Array.isArray(options?.apiVersion) ? options!.apiVersion @@ -300,7 +299,6 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { }, async list() { if (listActionsCache) return listActionsCache; - // Must return UNNAMESPACED names in v2. Genkit will prefix them. listActionsCache = await listActions(options); return listActionsCache; }, From 18f0670f11b9cc0ffe7d5646a7c1ae90724988ab Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Tue, 26 Aug 2025 15:06:52 +0100 Subject: [PATCH 4/7] refactor(plugins/googleai): revert adding prefixes --- js/plugins/googleai/src/embedder.ts | 18 +++++----- js/plugins/googleai/src/gemini.ts | 43 ++++++++++++------------ js/plugins/googleai/src/index.ts | 18 ++++++---- js/plugins/googleai/tests/gemini_test.ts | 1 + 4 files changed, 44 insertions(+), 36 deletions(-) diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index e1b219f7d6..6cda2e80b3 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -97,9 +97,9 @@ export const geminiEmbedding001 = embedderRef({ }); export const SUPPORTED_MODELS = { - 'googleai/embedding-001': textEmbeddingGecko001, - 'googleai/text-embedding-004': textEmbedding004, - 'googleai/gemini-embedding-001': geminiEmbedding001, + 'embedding-001': textEmbeddingGecko001, + 'text-embedding-004': textEmbedding004, + 'gemini-embedding-001': geminiEmbedding001, }; export function defineGoogleAIEmbedder( @@ -116,22 +116,24 @@ export function defineGoogleAIEmbedder( 'For more details see https://genkit.dev/docs/plugins/google-genai' ); } + // Extract the bare model name for lookup and API calls + const apiModelName = name.startsWith('googleai/') + ? name.substring('googleai/'.length) + : name; + const embedderReference: EmbedderReference = - SUPPORTED_MODELS[name] ?? + SUPPORTED_MODELS[apiModelName] ?? embedderRef({ name: name, configSchema: GeminiEmbeddingConfigSchema, info: { dimensions: 768, - label: `Google AI - ${name}`, + label: `Google AI - ${apiModelName}`, supports: { input: ['text', 'image', 'video'], }, }, }); - const apiModelName = embedderReference.name.startsWith('googleai/') - ? embedderReference.name.substring('googleai/'.length) - : embedderReference.name; return embedder( { name: embedderReference.name, diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index 500542fa65..db4ca1b6a0 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -59,7 +59,7 @@ import { import { downloadRequestMedia } from 'genkit/model/middleware'; import { model } from 'genkit/plugin'; import { runInNewSpan } from 'genkit/tracing'; -import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common.js'; +import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common'; import { handleCacheIfNeeded } from './context-caching'; import { extractCacheConfig } from './context-caching/utils'; @@ -645,26 +645,26 @@ export const gemma3ne4bit = modelRef({ }); export const SUPPORTED_GEMINI_MODELS = { - 'googleai/gemini-1.5-pro': gemini15Pro, - 'googleai/gemini-1.5-flash': gemini15Flash, - 'googleai/gemini-1.5-flash-8b': gemini15Flash8b, - 'googleai/gemini-2.0-pro-exp-02-05': gemini20ProExp0205, - 'googleai/gemini-2.0-flash': gemini20Flash, - 'googleai/gemini-2.0-flash-lite': gemini20FlashLite, - 'googleai/gemini-2.0-flash-exp': gemini20FlashExp, - 'googleai/gemini-2.5-pro-exp-03-25': gemini25ProExp0325, - 'googleai/gemini-2.5-pro-preview-03-25': gemini25ProPreview0325, - 'googleai/gemini-2.5-pro-preview-tts': gemini25ProPreviewTts, - 'googleai/gemini-2.5-flash-preview-04-17': gemini25FlashPreview0417, - 'googleai/gemini-2.5-flash-preview-tts': gemini25FlashPreviewTts, - 'googleai/gemini-2.5-flash': gemini25Flash, - 'googleai/gemini-2.5-flash-lite': gemini25FlashLite, - 'googleai/gemini-2.5-pro': gemini25Pro, - 'googleai/gemma-3-12b-it': gemma312bit, - 'googleai/gemma-3-1b-it': gemma31bit, - 'googleai/gemma-3-27b-it': gemma327bit, - 'googleai/gemma-3-4b-it': gemma34bit, - 'googleai/gemma-3n-e4b-it': gemma3ne4bit, + 'gemini-1.5-pro': gemini15Pro, + 'gemini-1.5-flash': gemini15Flash, + 'gemini-1.5-flash-8b': gemini15Flash8b, + 'gemini-2.0-pro-exp-02-05': gemini20ProExp0205, + 'gemini-2.0-flash': gemini20Flash, + 'gemini-2.0-flash-lite': gemini20FlashLite, + 'gemini-2.0-flash-exp': gemini20FlashExp, + 'gemini-2.5-pro-exp-03-25': gemini25ProExp0325, + 'gemini-2.5-pro-preview-03-25': gemini25ProPreview0325, + 'gemini-2.5-pro-preview-tts': gemini25ProPreviewTts, + 'gemini-2.5-flash-preview-04-17': gemini25FlashPreview0417, + 'gemini-2.5-flash-preview-tts': gemini25FlashPreviewTts, + 'gemini-2.5-flash': gemini25Flash, + 'gemini-2.5-flash-lite': gemini25FlashLite, + 'gemini-2.5-pro': gemini25Pro, + 'gemma-3-12b-it': gemma312bit, + 'gemma-3-1b-it': gemma31bit, + 'gemma-3-27b-it': gemma327bit, + 'gemma-3-4b-it': gemma34bit, + 'gemma-3n-e4b-it': gemma3ne4bit, }; export const GENERIC_GEMINI_MODEL = modelRef({ @@ -1149,6 +1149,7 @@ export function defineGoogleAIModel({ } } + // Extract the API model name for lookup and API calls const apiModelName = name.startsWith('googleai/') ? name.substring('googleai/'.length) : name; diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts index ca054ceba8..176f58853c 100644 --- a/js/plugins/googleai/src/index.ts +++ b/js/plugins/googleai/src/index.ts @@ -227,7 +227,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { actions.push( defineGoogleAIModel({ - name, + name: `googleai/${name}`, apiKey: options?.apiKey, apiVersion: 'v1beta', baseUrl: options?.baseUrl, @@ -241,7 +241,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { actions.push( defineGoogleAIModel({ - name, + name: `googleai/${name}`, apiKey: options?.apiKey, baseUrl: options?.baseUrl, debugTraces: options?.experimental_debugTraces, @@ -250,7 +250,9 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { }); Object.keys(EMBEDDER_MODELS).forEach((name) => { actions.push( - defineGoogleAIEmbedder(name, { apiKey: options?.apiKey }) + defineGoogleAIEmbedder(`googleai/${name}`, { + apiKey: options?.apiKey, + }) ); }); } @@ -278,18 +280,20 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { }, async resolve(actionType: ActionType, actionName: string) { if (actionType === 'embedder') { - return defineGoogleAIEmbedder(actionName, { apiKey: options?.apiKey }); + return defineGoogleAIEmbedder(`googleai/${actionName}`, { + apiKey: options?.apiKey, + }); } else if (actionName.startsWith('veo')) { if (actionType === 'background-model') { - return defineVeoModel(actionName, options?.apiKey); + return defineVeoModel(`googleai/${actionName}`, options?.apiKey); } } else if (actionType === 'model') { if (actionName.startsWith('imagen')) { - return defineImagenModel(actionName, options?.apiKey); + return defineImagenModel(`googleai/${actionName}`, options?.apiKey); } // Fallback dynamic Gemini model return defineGoogleAIModel({ - name: actionName, + name: `googleai/${actionName}`, apiKey: options?.apiKey, baseUrl: options?.baseUrl, debugTraces: options?.experimental_debugTraces, diff --git a/js/plugins/googleai/tests/gemini_test.ts b/js/plugins/googleai/tests/gemini_test.ts index 1a3cefa098..863ee3420b 100644 --- a/js/plugins/googleai/tests/gemini_test.ts +++ b/js/plugins/googleai/tests/gemini_test.ts @@ -619,6 +619,7 @@ describe('plugin', () => { toolChoice: true, systemRole: true, constrained: 'no-tools', + contextCache: true, }); const bananaRef = gemini('gemini-4.0-banana'); From f66475b213badea34512df4157eac836dfd39191 Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Thu, 28 Aug 2025 12:10:38 +0100 Subject: [PATCH 5/7] refactor(plugins/googleai): handle prefixes better --- js/plugins/googleai/src/common.ts | 29 ++ js/plugins/googleai/src/embedder.ts | 8 +- js/plugins/googleai/src/gemini.ts | 13 +- js/plugins/googleai/src/imagen.ts | 4 +- js/plugins/googleai/src/index.ts | 35 ++- js/plugins/googleai/src/veo.ts | 4 +- js/plugins/googleai/tests/common_test.ts | 338 +++++++++++++++++++++++ 7 files changed, 402 insertions(+), 29 deletions(-) create mode 100644 js/plugins/googleai/tests/common_test.ts diff --git a/js/plugins/googleai/src/common.ts b/js/plugins/googleai/src/common.ts index 4302a3c74d..3f87e3b470 100644 --- a/js/plugins/googleai/src/common.ts +++ b/js/plugins/googleai/src/common.ts @@ -31,3 +31,32 @@ export function getGenkitClientHeader() { } return defaultGetClientHeader(); } + +// Type-safe name helpers for GoogleAI plugin +const PROVIDER = 'googleai' as const; +type Provider = typeof PROVIDER; +type Prefixed = `${Provider}/${ActionName}`; +type MaybePrefixed = + | ActionName + | Prefixed; + +// Runtime + typed helpers +export function removePrefix( + name: MaybePrefixed +): ActionName { + return ( + name.startsWith(`${PROVIDER}/`) + ? (name.slice(PROVIDER.length + 1) as ActionName) + : name + ) as ActionName; +} + +export function ensurePrefixed( + name: MaybePrefixed +): Prefixed { + return ( + name.startsWith(`${PROVIDER}/`) + ? (name as Prefixed) + : (`${PROVIDER}/${name}` as Prefixed) + ) as Prefixed; +} diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index 6cda2e80b3..5916f44d64 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -26,8 +26,8 @@ import { } from 'genkit'; import { embedderRef } from 'genkit/embedder'; import { embedder } from 'genkit/plugin'; -import { getApiKeyFromEnvVar } from './common.js'; -import type { PluginOptions } from './index.js'; +import { getApiKeyFromEnvVar, removePrefix } from './common'; +import type { PluginOptions } from './index'; export const TaskTypeSchema = z.enum([ 'RETRIEVAL_DOCUMENT', @@ -117,9 +117,7 @@ export function defineGoogleAIEmbedder( ); } // Extract the bare model name for lookup and API calls - const apiModelName = name.startsWith('googleai/') - ? name.substring('googleai/'.length) - : name; + const apiModelName = removePrefix(name); const embedderReference: EmbedderReference = SUPPORTED_MODELS[apiModelName] ?? diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index db4ca1b6a0..7ca60044de 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -59,7 +59,12 @@ import { import { downloadRequestMedia } from 'genkit/model/middleware'; import { model } from 'genkit/plugin'; import { runInNewSpan } from 'genkit/tracing'; -import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common'; +import { + ensurePrefixed, + getApiKeyFromEnvVar, + getGenkitClientHeader, + removePrefix, +} from './common'; import { handleCacheIfNeeded } from './context-caching'; import { extractCacheConfig } from './context-caching/utils'; @@ -716,7 +721,7 @@ export function gemini( ): ModelReference { const nearestModel = nearestGeminiModelRef(version); return modelRef({ - name: `googleai/${version}`, + name: ensurePrefixed(version), config: options, configSchema: GeminiConfigSchema, info: { @@ -1150,9 +1155,7 @@ export function defineGoogleAIModel({ } // Extract the API model name for lookup and API calls - const apiModelName = name.startsWith('googleai/') - ? name.substring('googleai/'.length) - : name; + const apiModelName = removePrefix(name); const modelReference: ModelReference = SUPPORTED_GEMINI_MODELS[apiModelName] ?? diff --git a/js/plugins/googleai/src/imagen.ts b/js/plugins/googleai/src/imagen.ts index 335c4649aa..7004d4e058 100644 --- a/js/plugins/googleai/src/imagen.ts +++ b/js/plugins/googleai/src/imagen.ts @@ -24,7 +24,7 @@ import { type ModelReference, } from 'genkit/model'; import { model } from 'genkit/plugin'; -import { getApiKeyFromEnvVar } from './common.js'; +import { ensurePrefixed, getApiKeyFromEnvVar } from './common.js'; import { predictModel } from './predict.js'; export type KNOWN_IMAGEN_MODELS = 'imagen-3.0-generate-002'; @@ -124,7 +124,7 @@ export function defineImagenModel( }); } } - const modelName = `googleai/${name}`; + const modelName = ensurePrefixed(name); const modelReference: ModelReference = modelRef({ name: modelName, info: { diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts index 176f58853c..f832473398 100644 --- a/js/plugins/googleai/src/index.ts +++ b/js/plugins/googleai/src/index.ts @@ -27,7 +27,7 @@ import { logger } from 'genkit/logging'; import { modelRef } from 'genkit/model'; import { genkitPluginV2, type GenkitPluginV2 } from 'genkit/plugin'; import type { ActionType } from 'genkit/registry'; -import { getApiKeyFromEnvVar } from './common.js'; +import { ensurePrefixed, getApiKeyFromEnvVar, removePrefix } from './common'; import { SUPPORTED_MODELS as EMBEDDER_MODELS, GeminiEmbeddingConfigSchema, @@ -227,7 +227,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { actions.push( defineGoogleAIModel({ - name: `googleai/${name}`, + name: ensurePrefixed(name), apiKey: options?.apiKey, apiVersion: 'v1beta', baseUrl: options?.baseUrl, @@ -241,7 +241,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { actions.push( defineGoogleAIModel({ - name: `googleai/${name}`, + name: ensurePrefixed(name), apiKey: options?.apiKey, baseUrl: options?.baseUrl, debugTraces: options?.experimental_debugTraces, @@ -250,7 +250,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { }); Object.keys(EMBEDDER_MODELS).forEach((name) => { actions.push( - defineGoogleAIEmbedder(`googleai/${name}`, { + defineGoogleAIEmbedder(ensurePrefixed(name), { apiKey: options?.apiKey, }) ); @@ -267,7 +267,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef; actions.push( defineGoogleAIModel({ - name: modelName, + name: ensurePrefixed(modelName), apiKey: options?.apiKey, baseUrl: options?.baseUrl, info: { ...ref.info, label: `Google AI - ${modelName}` }, @@ -279,21 +279,24 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { return actions; }, async resolve(actionType: ActionType, actionName: string) { + const rawActionName = removePrefix(actionName); + const fullActionName = ensurePrefixed(rawActionName); + if (actionType === 'embedder') { - return defineGoogleAIEmbedder(`googleai/${actionName}`, { + return defineGoogleAIEmbedder(fullActionName, { apiKey: options?.apiKey, }); - } else if (actionName.startsWith('veo')) { + } else if (rawActionName.startsWith('veo')) { if (actionType === 'background-model') { - return defineVeoModel(`googleai/${actionName}`, options?.apiKey); + return defineVeoModel(fullActionName, options?.apiKey); } } else if (actionType === 'model') { - if (actionName.startsWith('imagen')) { - return defineImagenModel(`googleai/${actionName}`, options?.apiKey); + if (rawActionName.startsWith('imagen')) { + return defineImagenModel(fullActionName, options?.apiKey); } // Fallback dynamic Gemini model return defineGoogleAIModel({ - name: `googleai/${actionName}`, + name: fullActionName, apiKey: options?.apiKey, baseUrl: options?.baseUrl, debugTraces: options?.experimental_debugTraces, @@ -339,22 +342,23 @@ export const googleAI = googleAIPlugin as GoogleAIPlugin; name: string, config?: any ): ModelReference => { + const fullModelName = ensurePrefixed(name); if (name.startsWith('imagen')) { return modelRef({ - name: `googleai/${name}`, + name: fullModelName, config, configSchema: ImagenConfigSchema, }); } if (name.startsWith('veo')) { return modelRef({ - name: `googleai/${name}`, + name: fullModelName, config, configSchema: VeoConfigSchema, }); } return modelRef({ - name: `googleai/${name}`, + name: fullModelName, config, configSchema: GeminiConfigSchema, }); @@ -363,8 +367,9 @@ googleAI.embedder = ( name: string, config?: GeminiEmbeddingConfig ): EmbedderReference => { + const fullEmbedderName = ensurePrefixed(name); return embedderRef({ - name: `googleai/${name}`, + name: fullEmbedderName, config, configSchema: GeminiEmbeddingConfigSchema, }); diff --git a/js/plugins/googleai/src/veo.ts b/js/plugins/googleai/src/veo.ts index be1dba1bc7..ea4b1f0f13 100644 --- a/js/plugins/googleai/src/veo.ts +++ b/js/plugins/googleai/src/veo.ts @@ -23,7 +23,7 @@ import { type ModelReference, } from 'genkit/model'; import { backgroundModel } from 'genkit/plugin'; -import { getApiKeyFromEnvVar } from './common.js'; +import { ensurePrefixed, getApiKeyFromEnvVar } from './common.js'; import { Operation as ApiOperation, checkOp, predictModel } from './predict.js'; export type KNOWN_VEO_MODELS = 'veo-2.0-generate-001'; @@ -136,7 +136,7 @@ export function defineVeoModel( }); } } - const modelName = `googleai/${name}`; + const modelName = ensurePrefixed(name); const model: ModelReference = modelRef({ name: modelName, info: { diff --git a/js/plugins/googleai/tests/common_test.ts b/js/plugins/googleai/tests/common_test.ts new file mode 100644 index 0000000000..025f0e4ecc --- /dev/null +++ b/js/plugins/googleai/tests/common_test.ts @@ -0,0 +1,338 @@ +/** + * 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 assert from 'assert'; +import { describe, it } from 'node:test'; +import { ensurePrefixed, removePrefix } from '../src/common.js'; + +describe('Name Helper Functions', () => { + describe('removePrefix', () => { + it('should remove googleai/ prefix from prefixed names', () => { + assert.strictEqual( + removePrefix('googleai/gemini-1.5-flash'), + 'gemini-1.5-flash' + ); + assert.strictEqual( + removePrefix('googleai/imagen-3.0-generate-002'), + 'imagen-3.0-generate-002' + ); + assert.strictEqual( + removePrefix('googleai/veo-2.0-generate-001'), + 'veo-2.0-generate-001' + ); + assert.strictEqual( + removePrefix('googleai/embedding-001'), + 'embedding-001' + ); + }); + + it('should return unprefixed names unchanged', () => { + assert.strictEqual(removePrefix('gemini-1.5-flash'), 'gemini-1.5-flash'); + assert.strictEqual( + removePrefix('imagen-3.0-generate-002'), + 'imagen-3.0-generate-002' + ); + assert.strictEqual( + removePrefix('veo-2.0-generate-001'), + 'veo-2.0-generate-001' + ); + assert.strictEqual(removePrefix('embedding-001'), 'embedding-001'); + }); + + it('should handle edge cases', () => { + // Empty string + assert.strictEqual(removePrefix(''), ''); + + // Just the prefix + assert.strictEqual(removePrefix('googleai/'), ''); + + // Multiple slashes (should only remove the first googleai/) + assert.strictEqual( + removePrefix('googleai/googleai/gemini-1.5-flash'), + 'googleai/gemini-1.5-flash' + ); + + // Case sensitivity - should not match + assert.strictEqual( + removePrefix('GoogleAI/gemini-1.5-flash'), + 'GoogleAI/gemini-1.5-flash' + ); + assert.strictEqual( + removePrefix('GOOGLEAI/gemini-1.5-flash'), + 'GOOGLEAI/gemini-1.5-flash' + ); + }); + + it('should handle special characters in model names', () => { + assert.strictEqual( + removePrefix('googleai/gemini-2.0-flash-exp'), + 'gemini-2.0-flash-exp' + ); + assert.strictEqual( + removePrefix('googleai/gemini-2.5-pro-exp-03-25'), + 'gemini-2.5-pro-exp-03-25' + ); + assert.strictEqual( + removePrefix('googleai/gemma-3-12b-it'), + 'gemma-3-12b-it' + ); + }); + }); + + describe('ensurePrefixed', () => { + it('should add googleai/ prefix to unprefixed names', () => { + assert.strictEqual( + ensurePrefixed('gemini-1.5-flash'), + 'googleai/gemini-1.5-flash' + ); + assert.strictEqual( + ensurePrefixed('imagen-3.0-generate-002'), + 'googleai/imagen-3.0-generate-002' + ); + assert.strictEqual( + ensurePrefixed('veo-2.0-generate-001'), + 'googleai/veo-2.0-generate-001' + ); + assert.strictEqual( + ensurePrefixed('embedding-001'), + 'googleai/embedding-001' + ); + }); + + it('should return already prefixed names unchanged', () => { + assert.strictEqual( + ensurePrefixed('googleai/gemini-1.5-flash'), + 'googleai/gemini-1.5-flash' + ); + assert.strictEqual( + ensurePrefixed('googleai/imagen-3.0-generate-002'), + 'googleai/imagen-3.0-generate-002' + ); + assert.strictEqual( + ensurePrefixed('googleai/veo-2.0-generate-001'), + 'googleai/veo-2.0-generate-001' + ); + assert.strictEqual( + ensurePrefixed('googleai/embedding-001'), + 'googleai/embedding-001' + ); + }); + + it('should handle edge cases', () => { + // Empty string + assert.strictEqual(ensurePrefixed(''), 'googleai/'); + + // Just the prefix + assert.strictEqual(ensurePrefixed('googleai/'), 'googleai/'); + + // Multiple slashes (should add prefix to the whole thing) + assert.strictEqual( + ensurePrefixed('googleai/googleai/gemini-1.5-flash'), + 'googleai/googleai/gemini-1.5-flash' + ); + + // Case sensitivity - should not match existing prefix + assert.strictEqual( + ensurePrefixed('GoogleAI/gemini-1.5-flash'), + 'googleai/GoogleAI/gemini-1.5-flash' + ); + assert.strictEqual( + ensurePrefixed('GOOGLEAI/gemini-1.5-flash'), + 'googleai/GOOGLEAI/gemini-1.5-flash' + ); + }); + + it('should handle special characters in model names', () => { + assert.strictEqual( + ensurePrefixed('gemini-2.0-flash-exp'), + 'googleai/gemini-2.0-flash-exp' + ); + assert.strictEqual( + ensurePrefixed('gemini-2.5-pro-exp-03-25'), + 'googleai/gemini-2.5-pro-exp-03-25' + ); + assert.strictEqual( + ensurePrefixed('gemma-3-12b-it'), + 'googleai/gemma-3-12b-it' + ); + }); + }); + + describe('round-trip consistency', () => { + it('should maintain consistency: removePrefix(ensurePrefixed(x)) === x for non-prefix-only names', () => { + const testNames = [ + 'gemini-1.5-flash', + 'imagen-3.0-generate-002', + 'veo-2.0-generate-001', + 'embedding-001', + 'gemini-2.0-flash-exp', + 'gemini-2.5-pro-exp-03-25', + 'gemma-3-12b-it', + '', + ]; + + for (const name of testNames) { + const result = removePrefix(ensurePrefixed(name)); + assert.strictEqual(result, name, `Failed for name: "${name}"`); + } + }); + + it('should handle edge case: removePrefix(ensurePrefixed("googleai/")) === ""', () => { + // Special case: "googleai/" becomes "" after round-trip + const result = removePrefix(ensurePrefixed('googleai/')); + assert.strictEqual( + result, + '', + 'googleai/ should become empty string after round-trip' + ); + }); + + it('should maintain consistency: ensurePrefixed(removePrefix(x)) === x for prefixed names', () => { + const testNames = [ + 'googleai/gemini-1.5-flash', + 'googleai/imagen-3.0-generate-002', + 'googleai/veo-2.0-generate-001', + 'googleai/embedding-001', + 'googleai/gemini-2.0-flash-exp', + 'googleai/gemini-2.5-pro-exp-03-25', + 'googleai/gemma-3-12b-it', + 'googleai/', + ]; + + for (const name of testNames) { + const result = ensurePrefixed(removePrefix(name)); + assert.strictEqual(result, name, `Failed for name: "${name}"`); + } + }); + }); + + describe('double prefix prevention', () => { + it('should prevent double prefixes when used together', () => { + // Simulate the scenario that was causing double prefixes + const inputName = 'googleai/gemini-1.5-flash'; + + // This is what the resolver does: + const rawActionName = removePrefix(inputName); + const fullActionName = ensurePrefixed(rawActionName); + + // Should not result in double prefix + assert.strictEqual(fullActionName, 'googleai/gemini-1.5-flash'); + assert.notStrictEqual( + fullActionName, + 'googleai/googleai/gemini-1.5-flash' + ); + }); + + it('should handle multiple prefix removals safely', () => { + const doublePrefixed = 'googleai/googleai/gemini-1.5-flash'; + + // First removal + const firstRemoval = removePrefix(doublePrefixed); + assert.strictEqual(firstRemoval, 'googleai/gemini-1.5-flash'); + + // Second removal + const secondRemoval = removePrefix(firstRemoval); + assert.strictEqual(secondRemoval, 'gemini-1.5-flash'); + + // Ensure prefixed + const finalResult = ensurePrefixed(secondRemoval); + assert.strictEqual(finalResult, 'googleai/gemini-1.5-flash'); + }); + }); + + describe('TypeScript type safety', () => { + it('should work with template literal types', () => { + // These tests verify that the functions work with the template literal types + // defined in common.ts + + const unprefixedName = 'gemini-1.5-flash'; + const prefixedName = 'googleai/gemini-1.5-flash'; + + // removePrefix should return the raw name type + const rawResult = removePrefix(prefixedName); + assert.strictEqual(rawResult, unprefixedName); + + // ensurePrefixed should return the prefixed name type + const fullResult = ensurePrefixed(unprefixedName); + assert.strictEqual(fullResult, prefixedName); + + // Both should work with either input type + assert.strictEqual(removePrefix(unprefixedName), unprefixedName); + assert.strictEqual(ensurePrefixed(prefixedName), prefixedName); + }); + }); + + describe('real-world usage scenarios', () => { + it('should handle resolver scenarios correctly', () => { + // Simulate what happens in the plugin resolver + const scenarios = [ + { input: 'gemini-1.5-flash', expected: 'googleai/gemini-1.5-flash' }, + { + input: 'googleai/gemini-1.5-flash', + expected: 'googleai/gemini-1.5-flash', + }, + { + input: 'imagen-3.0-generate-002', + expected: 'googleai/imagen-3.0-generate-002', + }, + { + input: 'googleai/imagen-3.0-generate-002', + expected: 'googleai/imagen-3.0-generate-002', + }, + { + input: 'veo-2.0-generate-001', + expected: 'googleai/veo-2.0-generate-001', + }, + { + input: 'googleai/veo-2.0-generate-001', + expected: 'googleai/veo-2.0-generate-001', + }, + ]; + + for (const scenario of scenarios) { + const rawActionName = removePrefix(scenario.input); + const fullActionName = ensurePrefixed(rawActionName); + assert.strictEqual( + fullActionName, + scenario.expected, + `Failed for input: "${scenario.input}"` + ); + } + }); + + it('should handle static factory scenarios correctly', () => { + // Simulate what happens in the static .model() and .embedder() factories + const scenarios = [ + { input: 'gemini-1.5-flash', expected: 'googleai/gemini-1.5-flash' }, + { + input: 'googleai/gemini-1.5-flash', + expected: 'googleai/gemini-1.5-flash', + }, + { input: 'embedding-001', expected: 'googleai/embedding-001' }, + { input: 'googleai/embedding-001', expected: 'googleai/embedding-001' }, + ]; + + for (const scenario of scenarios) { + const fullName = ensurePrefixed(scenario.input); + assert.strictEqual( + fullName, + scenario.expected, + `Failed for input: "${scenario.input}"` + ); + } + }); + }); +}); From 19c542d76ac27ef36e7cb9bbc0bab70412ddd567 Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Thu, 28 Aug 2025 12:31:25 +0100 Subject: [PATCH 6/7] refactor(plugins/googleai): make prefix helpers generic --- js/plugins/googleai/src/common.ts | 88 ++++++++++++++++++++++------- js/plugins/googleai/src/embedder.ts | 4 +- js/plugins/googleai/src/gemini.ts | 2 +- 3 files changed, 71 insertions(+), 23 deletions(-) diff --git a/js/plugins/googleai/src/common.ts b/js/plugins/googleai/src/common.ts index 3f87e3b470..65da4a9305 100644 --- a/js/plugins/googleai/src/common.ts +++ b/js/plugins/googleai/src/common.ts @@ -32,31 +32,79 @@ export function getGenkitClientHeader() { return defaultGetClientHeader(); } -// Type-safe name helpers for GoogleAI plugin -const PROVIDER = 'googleai' as const; -type Provider = typeof PROVIDER; -type Prefixed = `${Provider}/${ActionName}`; -type MaybePrefixed = - | ActionName - | Prefixed; - -// Runtime + typed helpers -export function removePrefix( - name: MaybePrefixed +// Type-safe name helpers/guards + +const PREFIX = 'googleai' as const; + +type Prefix = typeof PREFIX; + +type Prefixed< + ActionName extends string, + PrefixType extends string = Prefix, +> = `${PrefixType}/${ActionName}`; + +type MaybePrefixed< + ActionName extends string, + PrefixType extends string = Prefix, +> = ActionName | Prefixed; + +/** + * Removes a prefix from an action name if it exists. + * + * @template ActionName - The action name type + * @template PrefixType - The prefix type (defaults to 'googleai') + * + * @param name - The action name, which may or may not be prefixed + * @param prefix - The prefix to remove (defaults to 'googleai') + * + * @returns The action name without the prefix + * + * @example + * ```typescript + * removePrefix('googleai/gemini-1.5-flash') // 'gemini-1.5-flash' + * removePrefix('gemini-1.5-flash') // 'gemini-1.5-flash' + * removePrefix('openai/gpt-4', 'openai') // 'gpt-4' + * ``` + */ +export function removePrefix< + ActionName extends string, + PrefixType extends string = Prefix, +>( + name: MaybePrefixed, + prefix: PrefixType = PREFIX as PrefixType ): ActionName { return ( - name.startsWith(`${PROVIDER}/`) - ? (name.slice(PROVIDER.length + 1) as ActionName) + name.startsWith(`${prefix}/`) + ? (name.slice(prefix.length + 1) as ActionName) : name ) as ActionName; } -export function ensurePrefixed( - name: MaybePrefixed -): Prefixed { +/** + * This function adds the prefix if it's missing, or returns the name unchanged + * if it already has the correct prefix. This prevents double-prefixing issues. + * + * @param name - The action name, which may or may not be prefixed + * @param prefix - The prefix to ensure (defaults to 'googleai') + * + * @returns The action name with the prefix guaranteed to be present + * + * @example + * ```typescript + * ensurePrefixed('gemini-1.5-flash') // 'googleai/gemini-1.5-flash' + * ensurePrefixed('googleai/gemini-1.5-flash') // 'googleai/gemini-1.5-flash' + * ``` + */ +export function ensurePrefixed< + ActionName extends string, + PrefixType extends string = Prefix, +>( + name: MaybePrefixed, + prefix: PrefixType = PREFIX as PrefixType +): Prefixed { return ( - name.startsWith(`${PROVIDER}/`) - ? (name as Prefixed) - : (`${PROVIDER}/${name}` as Prefixed) - ) as Prefixed; + name.startsWith(`${prefix}/`) + ? (name as Prefixed) + : (`${prefix}/${name}` as Prefixed) + ) as Prefixed; } diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index 5916f44d64..81e86ab741 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -26,7 +26,7 @@ import { } from 'genkit'; import { embedderRef } from 'genkit/embedder'; import { embedder } from 'genkit/plugin'; -import { getApiKeyFromEnvVar, removePrefix } from './common'; +import { ensurePrefixed, getApiKeyFromEnvVar, removePrefix } from './common'; import type { PluginOptions } from './index'; export const TaskTypeSchema = z.enum([ @@ -122,7 +122,7 @@ export function defineGoogleAIEmbedder( const embedderReference: EmbedderReference = SUPPORTED_MODELS[apiModelName] ?? embedderRef({ - name: name, + name: ensurePrefixed(name), configSchema: GeminiEmbeddingConfigSchema, info: { dimensions: 768, diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index 7ca60044de..cbe1290ded 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -1160,7 +1160,7 @@ export function defineGoogleAIModel({ const modelReference: ModelReference = SUPPORTED_GEMINI_MODELS[apiModelName] ?? modelRef({ - name: name, // Keep the full name for the model reference + name: ensurePrefixed(name), info: { label: `Google AI - ${apiModelName}`, supports: { From 92a62f00597b2d38795b60043e26d29b5bb0ba46 Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Tue, 9 Sep 2025 16:04:28 +0100 Subject: [PATCH 7/7] fix(plugins/googleai): revert prefixing code --- js/plugins/googleai/src/common.ts | 77 ------ js/plugins/googleai/src/embedder.ts | 16 +- js/plugins/googleai/src/gemini.ts | 25 +- js/plugins/googleai/src/imagen.ts | 13 +- js/plugins/googleai/src/index.ts | 36 ++- js/plugins/googleai/src/veo.ts | 13 +- js/plugins/googleai/tests/common_test.ts | 338 ----------------------- 7 files changed, 48 insertions(+), 470 deletions(-) delete mode 100644 js/plugins/googleai/tests/common_test.ts diff --git a/js/plugins/googleai/src/common.ts b/js/plugins/googleai/src/common.ts index 65da4a9305..4302a3c74d 100644 --- a/js/plugins/googleai/src/common.ts +++ b/js/plugins/googleai/src/common.ts @@ -31,80 +31,3 @@ export function getGenkitClientHeader() { } return defaultGetClientHeader(); } - -// Type-safe name helpers/guards - -const PREFIX = 'googleai' as const; - -type Prefix = typeof PREFIX; - -type Prefixed< - ActionName extends string, - PrefixType extends string = Prefix, -> = `${PrefixType}/${ActionName}`; - -type MaybePrefixed< - ActionName extends string, - PrefixType extends string = Prefix, -> = ActionName | Prefixed; - -/** - * Removes a prefix from an action name if it exists. - * - * @template ActionName - The action name type - * @template PrefixType - The prefix type (defaults to 'googleai') - * - * @param name - The action name, which may or may not be prefixed - * @param prefix - The prefix to remove (defaults to 'googleai') - * - * @returns The action name without the prefix - * - * @example - * ```typescript - * removePrefix('googleai/gemini-1.5-flash') // 'gemini-1.5-flash' - * removePrefix('gemini-1.5-flash') // 'gemini-1.5-flash' - * removePrefix('openai/gpt-4', 'openai') // 'gpt-4' - * ``` - */ -export function removePrefix< - ActionName extends string, - PrefixType extends string = Prefix, ->( - name: MaybePrefixed, - prefix: PrefixType = PREFIX as PrefixType -): ActionName { - return ( - name.startsWith(`${prefix}/`) - ? (name.slice(prefix.length + 1) as ActionName) - : name - ) as ActionName; -} - -/** - * This function adds the prefix if it's missing, or returns the name unchanged - * if it already has the correct prefix. This prevents double-prefixing issues. - * - * @param name - The action name, which may or may not be prefixed - * @param prefix - The prefix to ensure (defaults to 'googleai') - * - * @returns The action name with the prefix guaranteed to be present - * - * @example - * ```typescript - * ensurePrefixed('gemini-1.5-flash') // 'googleai/gemini-1.5-flash' - * ensurePrefixed('googleai/gemini-1.5-flash') // 'googleai/gemini-1.5-flash' - * ``` - */ -export function ensurePrefixed< - ActionName extends string, - PrefixType extends string = Prefix, ->( - name: MaybePrefixed, - prefix: PrefixType = PREFIX as PrefixType -): Prefixed { - return ( - name.startsWith(`${prefix}/`) - ? (name as Prefixed) - : (`${prefix}/${name}` as Prefixed) - ) as Prefixed; -} diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index 81e86ab741..4654b1f8f3 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -26,7 +26,7 @@ import { } from 'genkit'; import { embedderRef } from 'genkit/embedder'; import { embedder } from 'genkit/plugin'; -import { ensurePrefixed, getApiKeyFromEnvVar, removePrefix } from './common'; +import { getApiKeyFromEnvVar } from './common'; import type { PluginOptions } from './index'; export const TaskTypeSchema = z.enum([ @@ -116,17 +116,17 @@ export function defineGoogleAIEmbedder( 'For more details see https://genkit.dev/docs/plugins/google-genai' ); } - // Extract the bare model name for lookup and API calls - const apiModelName = removePrefix(name); + // In v2, plugin internals use UNPREFIXED action names. + const actionName = name; const embedderReference: EmbedderReference = - SUPPORTED_MODELS[apiModelName] ?? + SUPPORTED_MODELS[actionName] ?? embedderRef({ - name: ensurePrefixed(name), + name: actionName, configSchema: GeminiEmbeddingConfigSchema, info: { dimensions: 768, - label: `Google AI - ${apiModelName}`, + label: `Google AI - ${actionName}`, supports: { input: ['text', 'image', 'video'], }, @@ -134,7 +134,7 @@ export function defineGoogleAIEmbedder( }); return embedder( { - name: embedderReference.name, + name: actionName, configSchema: GeminiEmbeddingConfigSchema, info: embedderReference.info!, }, @@ -153,7 +153,7 @@ export function defineGoogleAIEmbedder( options?.version || embedderReference.config?.version || embedderReference.version || - apiModelName, + actionName, }); const embeddings = await Promise.all( input.map(async (doc) => { diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index cbe1290ded..7fd01cebd6 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -59,12 +59,7 @@ import { import { downloadRequestMedia } from 'genkit/model/middleware'; import { model } from 'genkit/plugin'; import { runInNewSpan } from 'genkit/tracing'; -import { - ensurePrefixed, - getApiKeyFromEnvVar, - getGenkitClientHeader, - removePrefix, -} from './common'; +import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common'; import { handleCacheIfNeeded } from './context-caching'; import { extractCacheConfig } from './context-caching/utils'; @@ -711,7 +706,7 @@ export type GeminiVersionString = * ```js * await ai.generate({ * prompt: 'hi', - * model: gemini('googleai/gemini-1.5-flash') + * model: gemini('gemini-1.5-flash') * }); * ``` */ @@ -721,7 +716,7 @@ export function gemini( ): ModelReference { const nearestModel = nearestGeminiModelRef(version); return modelRef({ - name: ensurePrefixed(version), + name: `googleai/${version}`, config: options, configSchema: GeminiConfigSchema, info: { @@ -1154,15 +1149,15 @@ export function defineGoogleAIModel({ } } - // Extract the API model name for lookup and API calls - const apiModelName = removePrefix(name); + // In v2, plugin internals use UNPREFIXED action names. + const actionName = name; const modelReference: ModelReference = - SUPPORTED_GEMINI_MODELS[apiModelName] ?? + SUPPORTED_GEMINI_MODELS[actionName] ?? modelRef({ - name: ensurePrefixed(name), + name: actionName, info: { - label: `Google AI - ${apiModelName}`, + label: `Google AI - ${actionName}`, supports: { multiturn: true, media: true, @@ -1204,7 +1199,7 @@ export function defineGoogleAIModel({ return model( { - name: modelReference.name, + name: actionName, ...modelReference.info, configSchema: modelReference.configSchema, use: middleware, @@ -1331,7 +1326,7 @@ export function defineGoogleAIModel({ } as StartChatParams; const modelVersion = (versionFromConfig || modelReference.version || - apiModelName) as string; + actionName) as string; const cacheConfigDetails = extractCacheConfig(request); const { chatRequest: updatedChatRequest, cache } = diff --git a/js/plugins/googleai/src/imagen.ts b/js/plugins/googleai/src/imagen.ts index 7004d4e058..433a9734a5 100644 --- a/js/plugins/googleai/src/imagen.ts +++ b/js/plugins/googleai/src/imagen.ts @@ -24,7 +24,7 @@ import { type ModelReference, } from 'genkit/model'; import { model } from 'genkit/plugin'; -import { ensurePrefixed, getApiKeyFromEnvVar } from './common.js'; +import { getApiKeyFromEnvVar } from './common.js'; import { predictModel } from './predict.js'; export type KNOWN_IMAGEN_MODELS = 'imagen-3.0-generate-002'; @@ -124,19 +124,20 @@ export function defineImagenModel( }); } } - const modelName = ensurePrefixed(name); + // In v2, plugin internals use UNPREFIXED action names. + const actionName = name; const modelReference: ModelReference = modelRef({ - name: modelName, + name: actionName, info: { ...GENERIC_IMAGEN_INFO, - label: `Google AI - ${name}`, + label: `Google AI - ${actionName}`, }, configSchema: ImagenConfigSchema, }); return model( { - name: modelName, + name: actionName, ...modelReference.info, configSchema: ImagenConfigSchema, }, @@ -153,7 +154,7 @@ export function defineImagenModel( ImagenInstance, ImagenPrediction, ImagenParameters - >(modelReference.version || name, apiKey as string, 'predict'); + >(modelReference.version || actionName, apiKey as string, 'predict'); const response = await predictClient([instance], toParameters(request)); if (!response.predictions || response.predictions.length == 0) { diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts index f832473398..e5caa55eda 100644 --- a/js/plugins/googleai/src/index.ts +++ b/js/plugins/googleai/src/index.ts @@ -27,7 +27,7 @@ import { logger } from 'genkit/logging'; import { modelRef } from 'genkit/model'; import { genkitPluginV2, type GenkitPluginV2 } from 'genkit/plugin'; import type { ActionType } from 'genkit/registry'; -import { ensurePrefixed, getApiKeyFromEnvVar, removePrefix } from './common'; +import { getApiKeyFromEnvVar } from './common'; import { SUPPORTED_MODELS as EMBEDDER_MODELS, GeminiEmbeddingConfigSchema, @@ -227,7 +227,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { actions.push( defineGoogleAIModel({ - name: ensurePrefixed(name), + name: name, apiKey: options?.apiKey, apiVersion: 'v1beta', baseUrl: options?.baseUrl, @@ -241,7 +241,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => { actions.push( defineGoogleAIModel({ - name: ensurePrefixed(name), + name: name, apiKey: options?.apiKey, baseUrl: options?.baseUrl, debugTraces: options?.experimental_debugTraces, @@ -250,7 +250,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { }); Object.keys(EMBEDDER_MODELS).forEach((name) => { actions.push( - defineGoogleAIEmbedder(ensurePrefixed(name), { + defineGoogleAIEmbedder(name, { apiKey: options?.apiKey, }) ); @@ -267,7 +267,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef; actions.push( defineGoogleAIModel({ - name: ensurePrefixed(modelName), + name: modelName, apiKey: options?.apiKey, baseUrl: options?.baseUrl, info: { ...ref.info, label: `Google AI - ${modelName}` }, @@ -279,24 +279,22 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { return actions; }, async resolve(actionType: ActionType, actionName: string) { - const rawActionName = removePrefix(actionName); - const fullActionName = ensurePrefixed(rawActionName); - + // In v2, actionName is already unprefixed if (actionType === 'embedder') { - return defineGoogleAIEmbedder(fullActionName, { + return defineGoogleAIEmbedder(actionName, { apiKey: options?.apiKey, }); - } else if (rawActionName.startsWith('veo')) { + } else if (actionName.startsWith('veo')) { if (actionType === 'background-model') { - return defineVeoModel(fullActionName, options?.apiKey); + return defineVeoModel(actionName, options?.apiKey); } } else if (actionType === 'model') { - if (rawActionName.startsWith('imagen')) { - return defineImagenModel(fullActionName, options?.apiKey); + if (actionName.startsWith('imagen')) { + return defineImagenModel(actionName, options?.apiKey); } // Fallback dynamic Gemini model return defineGoogleAIModel({ - name: fullActionName, + name: actionName, apiKey: options?.apiKey, baseUrl: options?.baseUrl, debugTraces: options?.experimental_debugTraces, @@ -342,23 +340,22 @@ export const googleAI = googleAIPlugin as GoogleAIPlugin; name: string, config?: any ): ModelReference => { - const fullModelName = ensurePrefixed(name); if (name.startsWith('imagen')) { return modelRef({ - name: fullModelName, + name: name, config, configSchema: ImagenConfigSchema, }); } if (name.startsWith('veo')) { return modelRef({ - name: fullModelName, + name: name, config, configSchema: VeoConfigSchema, }); } return modelRef({ - name: fullModelName, + name: name, config, configSchema: GeminiConfigSchema, }); @@ -367,9 +364,8 @@ googleAI.embedder = ( name: string, config?: GeminiEmbeddingConfig ): EmbedderReference => { - const fullEmbedderName = ensurePrefixed(name); return embedderRef({ - name: fullEmbedderName, + name: name, config, configSchema: GeminiEmbeddingConfigSchema, }); diff --git a/js/plugins/googleai/src/veo.ts b/js/plugins/googleai/src/veo.ts index ea4b1f0f13..20d4ed5e66 100644 --- a/js/plugins/googleai/src/veo.ts +++ b/js/plugins/googleai/src/veo.ts @@ -23,7 +23,7 @@ import { type ModelReference, } from 'genkit/model'; import { backgroundModel } from 'genkit/plugin'; -import { ensurePrefixed, getApiKeyFromEnvVar } from './common.js'; +import { getApiKeyFromEnvVar } from './common.js'; import { Operation as ApiOperation, checkOp, predictModel } from './predict.js'; export type KNOWN_VEO_MODELS = 'veo-2.0-generate-001'; @@ -136,18 +136,19 @@ export function defineVeoModel( }); } } - const modelName = ensurePrefixed(name); + // In v2, plugin internals use UNPREFIXED action names. + const actionName = name; const model: ModelReference = modelRef({ - name: modelName, + name: actionName, info: { ...GENERIC_VEO_INFO, - label: `Google AI - ${name}`, + label: `Google AI - ${actionName}`, }, configSchema: VeoConfigSchema, }); return backgroundModel({ - name: modelName, + name: actionName, ...model.info, configSchema: VeoConfigSchema, async start(request) { @@ -163,7 +164,7 @@ export function defineVeoModel( VeoInstance, ApiOperation, VeoParameters - >(model.version || name, apiKey as string, 'predictLongRunning'); + >(model.version || actionName, apiKey as string, 'predictLongRunning'); const response = await predictClient([instance], toParameters(request)); return toGenkitOp(response); diff --git a/js/plugins/googleai/tests/common_test.ts b/js/plugins/googleai/tests/common_test.ts deleted file mode 100644 index 025f0e4ecc..0000000000 --- a/js/plugins/googleai/tests/common_test.ts +++ /dev/null @@ -1,338 +0,0 @@ -/** - * 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 assert from 'assert'; -import { describe, it } from 'node:test'; -import { ensurePrefixed, removePrefix } from '../src/common.js'; - -describe('Name Helper Functions', () => { - describe('removePrefix', () => { - it('should remove googleai/ prefix from prefixed names', () => { - assert.strictEqual( - removePrefix('googleai/gemini-1.5-flash'), - 'gemini-1.5-flash' - ); - assert.strictEqual( - removePrefix('googleai/imagen-3.0-generate-002'), - 'imagen-3.0-generate-002' - ); - assert.strictEqual( - removePrefix('googleai/veo-2.0-generate-001'), - 'veo-2.0-generate-001' - ); - assert.strictEqual( - removePrefix('googleai/embedding-001'), - 'embedding-001' - ); - }); - - it('should return unprefixed names unchanged', () => { - assert.strictEqual(removePrefix('gemini-1.5-flash'), 'gemini-1.5-flash'); - assert.strictEqual( - removePrefix('imagen-3.0-generate-002'), - 'imagen-3.0-generate-002' - ); - assert.strictEqual( - removePrefix('veo-2.0-generate-001'), - 'veo-2.0-generate-001' - ); - assert.strictEqual(removePrefix('embedding-001'), 'embedding-001'); - }); - - it('should handle edge cases', () => { - // Empty string - assert.strictEqual(removePrefix(''), ''); - - // Just the prefix - assert.strictEqual(removePrefix('googleai/'), ''); - - // Multiple slashes (should only remove the first googleai/) - assert.strictEqual( - removePrefix('googleai/googleai/gemini-1.5-flash'), - 'googleai/gemini-1.5-flash' - ); - - // Case sensitivity - should not match - assert.strictEqual( - removePrefix('GoogleAI/gemini-1.5-flash'), - 'GoogleAI/gemini-1.5-flash' - ); - assert.strictEqual( - removePrefix('GOOGLEAI/gemini-1.5-flash'), - 'GOOGLEAI/gemini-1.5-flash' - ); - }); - - it('should handle special characters in model names', () => { - assert.strictEqual( - removePrefix('googleai/gemini-2.0-flash-exp'), - 'gemini-2.0-flash-exp' - ); - assert.strictEqual( - removePrefix('googleai/gemini-2.5-pro-exp-03-25'), - 'gemini-2.5-pro-exp-03-25' - ); - assert.strictEqual( - removePrefix('googleai/gemma-3-12b-it'), - 'gemma-3-12b-it' - ); - }); - }); - - describe('ensurePrefixed', () => { - it('should add googleai/ prefix to unprefixed names', () => { - assert.strictEqual( - ensurePrefixed('gemini-1.5-flash'), - 'googleai/gemini-1.5-flash' - ); - assert.strictEqual( - ensurePrefixed('imagen-3.0-generate-002'), - 'googleai/imagen-3.0-generate-002' - ); - assert.strictEqual( - ensurePrefixed('veo-2.0-generate-001'), - 'googleai/veo-2.0-generate-001' - ); - assert.strictEqual( - ensurePrefixed('embedding-001'), - 'googleai/embedding-001' - ); - }); - - it('should return already prefixed names unchanged', () => { - assert.strictEqual( - ensurePrefixed('googleai/gemini-1.5-flash'), - 'googleai/gemini-1.5-flash' - ); - assert.strictEqual( - ensurePrefixed('googleai/imagen-3.0-generate-002'), - 'googleai/imagen-3.0-generate-002' - ); - assert.strictEqual( - ensurePrefixed('googleai/veo-2.0-generate-001'), - 'googleai/veo-2.0-generate-001' - ); - assert.strictEqual( - ensurePrefixed('googleai/embedding-001'), - 'googleai/embedding-001' - ); - }); - - it('should handle edge cases', () => { - // Empty string - assert.strictEqual(ensurePrefixed(''), 'googleai/'); - - // Just the prefix - assert.strictEqual(ensurePrefixed('googleai/'), 'googleai/'); - - // Multiple slashes (should add prefix to the whole thing) - assert.strictEqual( - ensurePrefixed('googleai/googleai/gemini-1.5-flash'), - 'googleai/googleai/gemini-1.5-flash' - ); - - // Case sensitivity - should not match existing prefix - assert.strictEqual( - ensurePrefixed('GoogleAI/gemini-1.5-flash'), - 'googleai/GoogleAI/gemini-1.5-flash' - ); - assert.strictEqual( - ensurePrefixed('GOOGLEAI/gemini-1.5-flash'), - 'googleai/GOOGLEAI/gemini-1.5-flash' - ); - }); - - it('should handle special characters in model names', () => { - assert.strictEqual( - ensurePrefixed('gemini-2.0-flash-exp'), - 'googleai/gemini-2.0-flash-exp' - ); - assert.strictEqual( - ensurePrefixed('gemini-2.5-pro-exp-03-25'), - 'googleai/gemini-2.5-pro-exp-03-25' - ); - assert.strictEqual( - ensurePrefixed('gemma-3-12b-it'), - 'googleai/gemma-3-12b-it' - ); - }); - }); - - describe('round-trip consistency', () => { - it('should maintain consistency: removePrefix(ensurePrefixed(x)) === x for non-prefix-only names', () => { - const testNames = [ - 'gemini-1.5-flash', - 'imagen-3.0-generate-002', - 'veo-2.0-generate-001', - 'embedding-001', - 'gemini-2.0-flash-exp', - 'gemini-2.5-pro-exp-03-25', - 'gemma-3-12b-it', - '', - ]; - - for (const name of testNames) { - const result = removePrefix(ensurePrefixed(name)); - assert.strictEqual(result, name, `Failed for name: "${name}"`); - } - }); - - it('should handle edge case: removePrefix(ensurePrefixed("googleai/")) === ""', () => { - // Special case: "googleai/" becomes "" after round-trip - const result = removePrefix(ensurePrefixed('googleai/')); - assert.strictEqual( - result, - '', - 'googleai/ should become empty string after round-trip' - ); - }); - - it('should maintain consistency: ensurePrefixed(removePrefix(x)) === x for prefixed names', () => { - const testNames = [ - 'googleai/gemini-1.5-flash', - 'googleai/imagen-3.0-generate-002', - 'googleai/veo-2.0-generate-001', - 'googleai/embedding-001', - 'googleai/gemini-2.0-flash-exp', - 'googleai/gemini-2.5-pro-exp-03-25', - 'googleai/gemma-3-12b-it', - 'googleai/', - ]; - - for (const name of testNames) { - const result = ensurePrefixed(removePrefix(name)); - assert.strictEqual(result, name, `Failed for name: "${name}"`); - } - }); - }); - - describe('double prefix prevention', () => { - it('should prevent double prefixes when used together', () => { - // Simulate the scenario that was causing double prefixes - const inputName = 'googleai/gemini-1.5-flash'; - - // This is what the resolver does: - const rawActionName = removePrefix(inputName); - const fullActionName = ensurePrefixed(rawActionName); - - // Should not result in double prefix - assert.strictEqual(fullActionName, 'googleai/gemini-1.5-flash'); - assert.notStrictEqual( - fullActionName, - 'googleai/googleai/gemini-1.5-flash' - ); - }); - - it('should handle multiple prefix removals safely', () => { - const doublePrefixed = 'googleai/googleai/gemini-1.5-flash'; - - // First removal - const firstRemoval = removePrefix(doublePrefixed); - assert.strictEqual(firstRemoval, 'googleai/gemini-1.5-flash'); - - // Second removal - const secondRemoval = removePrefix(firstRemoval); - assert.strictEqual(secondRemoval, 'gemini-1.5-flash'); - - // Ensure prefixed - const finalResult = ensurePrefixed(secondRemoval); - assert.strictEqual(finalResult, 'googleai/gemini-1.5-flash'); - }); - }); - - describe('TypeScript type safety', () => { - it('should work with template literal types', () => { - // These tests verify that the functions work with the template literal types - // defined in common.ts - - const unprefixedName = 'gemini-1.5-flash'; - const prefixedName = 'googleai/gemini-1.5-flash'; - - // removePrefix should return the raw name type - const rawResult = removePrefix(prefixedName); - assert.strictEqual(rawResult, unprefixedName); - - // ensurePrefixed should return the prefixed name type - const fullResult = ensurePrefixed(unprefixedName); - assert.strictEqual(fullResult, prefixedName); - - // Both should work with either input type - assert.strictEqual(removePrefix(unprefixedName), unprefixedName); - assert.strictEqual(ensurePrefixed(prefixedName), prefixedName); - }); - }); - - describe('real-world usage scenarios', () => { - it('should handle resolver scenarios correctly', () => { - // Simulate what happens in the plugin resolver - const scenarios = [ - { input: 'gemini-1.5-flash', expected: 'googleai/gemini-1.5-flash' }, - { - input: 'googleai/gemini-1.5-flash', - expected: 'googleai/gemini-1.5-flash', - }, - { - input: 'imagen-3.0-generate-002', - expected: 'googleai/imagen-3.0-generate-002', - }, - { - input: 'googleai/imagen-3.0-generate-002', - expected: 'googleai/imagen-3.0-generate-002', - }, - { - input: 'veo-2.0-generate-001', - expected: 'googleai/veo-2.0-generate-001', - }, - { - input: 'googleai/veo-2.0-generate-001', - expected: 'googleai/veo-2.0-generate-001', - }, - ]; - - for (const scenario of scenarios) { - const rawActionName = removePrefix(scenario.input); - const fullActionName = ensurePrefixed(rawActionName); - assert.strictEqual( - fullActionName, - scenario.expected, - `Failed for input: "${scenario.input}"` - ); - } - }); - - it('should handle static factory scenarios correctly', () => { - // Simulate what happens in the static .model() and .embedder() factories - const scenarios = [ - { input: 'gemini-1.5-flash', expected: 'googleai/gemini-1.5-flash' }, - { - input: 'googleai/gemini-1.5-flash', - expected: 'googleai/gemini-1.5-flash', - }, - { input: 'embedding-001', expected: 'googleai/embedding-001' }, - { input: 'googleai/embedding-001', expected: 'googleai/embedding-001' }, - ]; - - for (const scenario of scenarios) { - const fullName = ensurePrefixed(scenario.input); - assert.strictEqual( - fullName, - scenario.expected, - `Failed for input: "${scenario.input}"` - ); - } - }); - }); -});