From 35e23184057cd40a36a5adada9867f76be4c511b Mon Sep 17 00:00:00 2001 From: Corie Watson Date: Mon, 27 Oct 2025 11:24:00 +0000 Subject: [PATCH 1/3] refactor(plugins/googleai): migrate to v2 --- js/plugins/googleai/src/embedder.ts | 33 +++++------ js/plugins/googleai/src/gemini.ts | 90 ++++++++++++++--------------- js/plugins/googleai/src/imagen.ts | 14 ++--- js/plugins/googleai/src/index.ts | 67 ++++++++++----------- js/plugins/googleai/src/veo.ts | 13 ++--- 5 files changed, 106 insertions(+), 111 deletions(-) diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index 531ac8a0e5..ea90c51c40 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -23,11 +23,11 @@ import { z, type EmbedderAction, type EmbedderReference, - type Genkit, } from 'genkit'; -import { embedderRef } from 'genkit/embedder'; +import { embedderRef as createEmbedderRef } from 'genkit/embedder'; import { getApiKeyFromEnvVar } from './common.js'; import type { PluginOptions } from './index.js'; +import { embedder } from 'genkit/plugin'; export const TaskTypeSchema = z.enum([ 'RETRIEVAL_DOCUMENT', @@ -60,7 +60,7 @@ export const GeminiEmbeddingConfigSchema = z.object({ export type GeminiEmbeddingConfig = z.infer; -export const textEmbeddingGecko001 = embedderRef({ +export const textEmbeddingGecko001 = createEmbedderRef({ name: 'googleai/embedding-001', configSchema: GeminiEmbeddingConfigSchema, info: { @@ -72,7 +72,7 @@ export const textEmbeddingGecko001 = embedderRef({ }, }); -export const textEmbedding004 = embedderRef({ +export const textEmbedding004 = createEmbedderRef({ name: 'googleai/text-embedding-004', configSchema: GeminiEmbeddingConfigSchema, info: { @@ -84,7 +84,7 @@ export const textEmbedding004 = embedderRef({ }, }); -export const geminiEmbedding001 = embedderRef({ +export const geminiEmbedding001 = createEmbedderRef({ name: 'googleai/gemini-embedding-001', configSchema: GeminiEmbeddingConfigSchema, info: { @@ -103,7 +103,6 @@ export const SUPPORTED_MODELS = { }; export function defineGoogleAIEmbedder( - ai: Genkit, name: string, pluginOptions: PluginOptions ): EmbedderAction { @@ -117,9 +116,9 @@ export function defineGoogleAIEmbedder( 'For more details see https://genkit.dev/docs/plugins/google-genai' ); } - const embedder: EmbedderReference = + const embedderRef: EmbedderReference = SUPPORTED_MODELS[name] ?? - embedderRef({ + createEmbedderRef({ name: name, configSchema: GeminiEmbeddingConfigSchema, info: { @@ -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 = embedderRef.name.startsWith('googleai/') + ? embedderRef.name.substring('googleai/'.length) + : embedderRef.name; + return embedder( { - name: embedder.name, + name: embedderRef.name, configSchema: GeminiEmbeddingConfigSchema, - info: embedder.info!, + info: embedderRef.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 || + embedderRef.config?.version || + embedderRef.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..cb81227078 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -38,12 +38,12 @@ 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, getBasicUsageStats, - modelRef, + modelRef as createModelRef, type CandidateData, type MediaPart, type MessageData, @@ -61,6 +61,7 @@ import { runInNewSpan } from 'genkit/tracing'; import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common'; import { handleCacheIfNeeded } from './context-caching'; import { extractCacheConfig } from './context-caching/utils'; +import { model } from 'genkit/plugin'; // Extra type guard to keep the compiler happy and avoid a cast to any. The // legacy Gemini SDK is no longer maintained, and doesn't have updated types. @@ -276,7 +277,7 @@ export const GeminiTtsConfigSchema = GeminiConfigSchema.extend({ .optional(), }).passthrough(); -export const gemini10Pro = modelRef({ +export const gemini10Pro = createModelRef({ name: 'googleai/gemini-1.0-pro', info: { label: 'Google AI - Gemini Pro', @@ -293,7 +294,7 @@ export const gemini10Pro = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini15Pro = modelRef({ +export const gemini15Pro = createModelRef({ name: 'googleai/gemini-1.5-pro', info: { label: 'Google AI - Gemini 1.5 Pro', @@ -314,7 +315,7 @@ export const gemini15Pro = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini15Flash = modelRef({ +export const gemini15Flash = createModelRef({ name: 'googleai/gemini-1.5-flash', info: { label: 'Google AI - Gemini 1.5 Flash', @@ -337,7 +338,7 @@ export const gemini15Flash = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini15Flash8b = modelRef({ +export const gemini15Flash8b = createModelRef({ name: 'googleai/gemini-1.5-flash-8b', info: { label: 'Google AI - Gemini 1.5 Flash', @@ -354,7 +355,7 @@ export const gemini15Flash8b = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini20Flash = modelRef({ +export const gemini20Flash = createModelRef({ name: 'googleai/gemini-2.0-flash', info: { label: 'Google AI - Gemini 2.0 Flash', @@ -371,7 +372,7 @@ export const gemini20Flash = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini20FlashExp = modelRef({ +export const gemini20FlashExp = createModelRef({ name: 'googleai/gemini-2.0-flash-exp', info: { label: 'Google AI - Gemini 2.0 Flash (Experimental)', @@ -388,7 +389,7 @@ export const gemini20FlashExp = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini20FlashLite = modelRef({ +export const gemini20FlashLite = createModelRef({ name: 'googleai/gemini-2.0-flash-lite', info: { label: 'Google AI - Gemini 2.0 Flash Lite', @@ -405,7 +406,7 @@ export const gemini20FlashLite = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini20ProExp0205 = modelRef({ +export const gemini20ProExp0205 = createModelRef({ name: 'googleai/gemini-2.0-pro-exp-02-05', info: { label: 'Google AI - Gemini 2.0 Pro Exp 02-05', @@ -422,7 +423,7 @@ export const gemini20ProExp0205 = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini25FlashPreview0417 = modelRef({ +export const gemini25FlashPreview0417 = createModelRef({ name: 'googleai/gemini-2.5-flash-preview-04-17', info: { label: 'Google AI - Gemini 2.5 Flash Preview 04-17', @@ -439,7 +440,7 @@ export const gemini25FlashPreview0417 = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini25FlashPreviewTts = modelRef({ +export const gemini25FlashPreviewTts = createModelRef({ name: 'googleai/gemini-2.5-flash-preview-tts', info: { label: 'Google AI - Gemini 2.5 Flash Preview TTS', @@ -456,7 +457,7 @@ export const gemini25FlashPreviewTts = modelRef({ configSchema: GeminiTtsConfigSchema, }); -export const gemini25ProExp0325 = modelRef({ +export const gemini25ProExp0325 = createModelRef({ name: 'googleai/gemini-2.5-pro-exp-03-25', info: { label: 'Google AI - Gemini 2.5 Pro Exp 03-25', @@ -473,7 +474,7 @@ export const gemini25ProExp0325 = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini25ProPreview0325 = modelRef({ +export const gemini25ProPreview0325 = createModelRef({ name: 'googleai/gemini-2.5-pro-preview-03-25', info: { label: 'Google AI - Gemini 2.5 Pro Preview 03-25', @@ -490,7 +491,7 @@ export const gemini25ProPreview0325 = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini25ProPreviewTts = modelRef({ +export const gemini25ProPreviewTts = createModelRef({ name: 'googleai/gemini-2.5-pro-preview-tts', info: { label: 'Google AI - Gemini 2.5 Pro Preview TTS', @@ -507,7 +508,7 @@ export const gemini25ProPreviewTts = modelRef({ configSchema: GeminiTtsConfigSchema, }); -export const gemini25Pro = modelRef({ +export const gemini25Pro = createModelRef({ name: 'googleai/gemini-2.5-pro', info: { label: 'Google AI - Gemini 2.5 Pro', @@ -524,7 +525,7 @@ export const gemini25Pro = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini25Flash = modelRef({ +export const gemini25Flash = createModelRef({ name: 'googleai/gemini-2.5-flash', info: { label: 'Google AI - Gemini 2.5 Flash', @@ -541,7 +542,7 @@ export const gemini25Flash = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemini25FlashLite = modelRef({ +export const gemini25FlashLite = createModelRef({ name: 'googleai/gemini-2.5-flash-lite', info: { label: 'Google AI - Gemini 2.5 Flash Lite', @@ -558,7 +559,7 @@ export const gemini25FlashLite = modelRef({ configSchema: GeminiConfigSchema, }); -export const gemma312bit = modelRef({ +export const gemma312bit = createModelRef({ name: 'googleai/gemma-3-12b-it', info: { label: 'Google AI - Gemma 3 12B', @@ -575,7 +576,7 @@ export const gemma312bit = modelRef({ configSchema: GeminiGemmaConfigSchema, }); -export const gemma31bit = modelRef({ +export const gemma31bit = createModelRef({ name: 'googleai/gemma-3-1b-it', info: { label: 'Google AI - Gemma 3 1B', @@ -592,7 +593,7 @@ export const gemma31bit = modelRef({ configSchema: GeminiGemmaConfigSchema, }); -export const gemma327bit = modelRef({ +export const gemma327bit = createModelRef({ name: 'googleai/gemma-3-27b-it', info: { label: 'Google AI - Gemma 3 27B', @@ -609,7 +610,7 @@ export const gemma327bit = modelRef({ configSchema: GeminiGemmaConfigSchema, }); -export const gemma34bit = modelRef({ +export const gemma34bit = createModelRef({ name: 'googleai/gemma-3-4b-it', info: { label: 'Google AI - Gemma 3 4B', @@ -626,7 +627,7 @@ export const gemma34bit = modelRef({ configSchema: GeminiGemmaConfigSchema, }); -export const gemma3ne4bit = modelRef({ +export const gemma3ne4bit = createModelRef({ name: 'googleai/gemma-3n-e4b-it', info: { label: 'Google AI - Gemma 3n E4B', @@ -666,7 +667,7 @@ export const SUPPORTED_GEMINI_MODELS = { 'gemma-3n-e4b-it': gemma3ne4bit, }; -export const GENERIC_GEMINI_MODEL = modelRef({ +export const GENERIC_GEMINI_MODEL = createModelRef({ name: 'googleai/gemini', configSchema: GeminiConfigSchema, info: { @@ -714,7 +715,7 @@ export function gemini( options: GeminiConfig = {} ): ModelReference { const nearestModel = nearestGeminiModelRef(version); - return modelRef({ + return createModelRef({ name: `googleai/${version}`, config: options, configSchema: GeminiConfigSchema, @@ -1118,7 +1119,6 @@ export function cleanSchema(schema: JSONSchema): JSONSchema { * Defines a new GoogleAI model. */ export function defineGoogleAIModel({ - ai, name, apiKey: apiKeyOption, apiVersion, @@ -1127,7 +1127,6 @@ export function defineGoogleAIModel({ defaultConfig, debugTraces, }: { - ai: Genkit; name: string; apiKey?: string | false; apiVersion?: string; @@ -1154,9 +1153,9 @@ export function defineGoogleAIModel({ ? name.substring('googleai/'.length) : name; - const model: ModelReference = + const modelRef: ModelReference = SUPPORTED_GEMINI_MODELS[apiModelName] ?? - modelRef({ + createModelRef({ name: `googleai/${apiModelName}`, info: { label: `Google AI - ${apiModelName}`, @@ -1173,7 +1172,7 @@ export function defineGoogleAIModel({ }); const middleware: ModelMiddleware[] = []; - if (model.info?.supports?.media) { + if (modelRef.info?.supports?.media) { // the gemini api doesn't support downloading media from http(s) middleware.push( downloadRequestMedia({ @@ -1199,15 +1198,14 @@ export function defineGoogleAIModel({ ); } - return ai.defineModel( + return model( { - apiVersion: 'v2', - name: model.name, - ...model.info, - configSchema: model.configSchema, + name: modelRef.name, + ...modelRef.info, + configSchema: GeminiConfigSchema, use: middleware, }, - async (request, { streamingRequested, sendChunk, abortSignal }) => { + async (request, { streamingRequested, sendChunk, abortSignal, registry }) => { const options: RequestOptions = { apiClient: getGenkitClientHeader() }; if (apiVersion) { options.apiVersion = apiVersion; @@ -1228,7 +1226,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 (modelRef.info?.supports?.systemRole) { const systemMessage = messages.find((m) => m.role === 'system'); if (systemMessage) { messages.splice(messages.indexOf(systemMessage), 1); @@ -1257,9 +1255,11 @@ export function defineGoogleAIModel({ if (codeExecutionFromConfig) { tools.push({ codeExecution: - request.config.codeExecution === true - ? {} - : request.config.codeExecution, + requestConfig.codeExecution ? + requestConfig.codeExecution === true + ? {} + : requestConfig.codeExecution + : {}, }); } @@ -1306,7 +1306,7 @@ export function defineGoogleAIModel({ generationConfig.responseSchema = cleanSchema(request.output.schema); } - const msg = toGeminiMessage(messages[messages.length - 1], model); + const msg = toGeminiMessage(messages[messages.length - 1], modelRef); 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, modelRef)), safetySettings: safetySettingsFromConfig, } as StartChatParams; const modelVersion = (versionFromConfig || - model.version || + modelRef.version || apiModelName) as string; const cacheConfigDetails = extractCacheConfig(request); @@ -1428,9 +1428,9 @@ 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 + return debugTraces && registry ? await runInNewSpan( - ai.registry, + registry, { metadata: { name: streamingRequested ? 'sendMessageStream' : 'sendMessage', diff --git a/js/plugins/googleai/src/imagen.ts b/js/plugins/googleai/src/imagen.ts index d7d6c39709..5209f42be3 100644 --- a/js/plugins/googleai/src/imagen.ts +++ b/js/plugins/googleai/src/imagen.ts @@ -14,10 +14,10 @@ * limitations under the License. */ -import { GenkitError, MessageData, z, type Genkit } from 'genkit'; +import { GenkitError, MessageData, z } from 'genkit'; import { getBasicUsageStats, - modelRef, + modelRef as createModelRef, type GenerateRequest, type ModelAction, type ModelInfo, @@ -25,6 +25,7 @@ import { } from 'genkit/model'; import { getApiKeyFromEnvVar } from './common.js'; import { predictModel } from './predict.js'; +import { model } from 'genkit/plugin'; export type KNOWN_IMAGEN_MODELS = 'imagen-3.0-generate-002'; @@ -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 modelRef: ModelReference = createModelRef({ name: modelName, info: { ...GENERIC_IMAGEN_INFO, @@ -134,10 +134,10 @@ export function defineImagenModel( configSchema: ImagenConfigSchema, }); - return ai.defineModel( + return model( { name: modelName, - ...model.info, + ...modelRef.info, configSchema: ImagenConfigSchema, }, async (request) => { @@ -153,7 +153,7 @@ export function defineImagenModel( ImagenInstance, ImagenPrediction, ImagenParameters - >(model.version || name, apiKey as string, 'predict'); + >(modelRef.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..092c6a81c1 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, ResolvableAction, type GenkitPluginV2 } from 'genkit/plugin'; import type { ActionType } from 'genkit/registry'; import { getApiKeyFromEnvVar } from './common.js'; import { @@ -112,8 +111,9 @@ export interface PluginOptions { experimental_debugTraces?: boolean; } -async function initializer(ai: Genkit, options?: PluginOptions) { +async function initializer(options?: PluginOptions): Promise { let apiVersions = ['v1']; + const actions: ResolvableAction[] = []; if (options?.apiVersion) { if (Array.isArray(options?.apiVersion)) { @@ -125,29 +125,27 @@ async function initializer(ai: Genkit, options?: PluginOptions) { if (apiVersions.includes('v1beta')) { Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => - defineGoogleAIModel({ - ai, + 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) => - defineGoogleAIModel({ - ai, + actions.push(defineGoogleAIModel({ 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 }) + actions.push(defineGoogleAIEmbedder(name, { apiKey: options?.apiKey })) ); } @@ -160,8 +158,7 @@ async function initializer(ai: Genkit, options?: PluginOptions) { modelOrRef.name.split('/')[1]; const modelRef = typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef; - defineGoogleAIModel({ - ai, + actions.push(defineGoogleAIModel({ name: modelName, apiKey: options?.apiKey, baseUrl: options?.baseUrl, @@ -170,40 +167,41 @@ async function initializer(ai: Genkit, options?: PluginOptions) { label: `Google AI - ${modelName}`, }, debugTraces: options?.experimental_debugTraces, - }); + })); } } + + return actions; } async function resolver( - ai: Genkit, actionType: ActionType, actionName: string, options?: PluginOptions -) { +): Promise { if (actionType === 'embedder') { - resolveEmbedder(ai, actionName, options); + return resolveEmbedder(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); + return defineVeoModel(actionName, options?.apiKey); } } else if (actionType === 'model') { - resolveModel(ai, actionName, options); + return resolveModel(actionName, options); } + + return undefined; } -function resolveModel(ai: Genkit, actionName: string, options?: PluginOptions) { +function resolveModel(actionName: string, options?: PluginOptions): ResolvableAction | undefined { if (actionName.startsWith('imagen')) { - defineImagenModel(ai, actionName, options?.apiKey); - return; + return defineImagenModel(actionName, options?.apiKey); } const modelRef = gemini(actionName); - defineGoogleAIModel({ - ai, + return defineGoogleAIModel({ name: modelRef.name, apiKey: options?.apiKey, baseUrl: options?.baseUrl, @@ -216,11 +214,10 @@ function resolveModel(ai: Genkit, actionName: string, options?: PluginOptions) { } function resolveEmbedder( - ai: Genkit, actionName: string, options?: PluginOptions -) { - defineGoogleAIEmbedder(ai, `googleai/${actionName}`, { +): ResolvableAction | undefined { + return defineGoogleAIEmbedder(`googleai/${actionName}`, { apiKey: options?.apiKey, }); } @@ -328,23 +325,23 @@ 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', + 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 = { - (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..9049701a37 100644 --- a/js/plugins/googleai/src/veo.ts +++ b/js/plugins/googleai/src/veo.ts @@ -19,17 +19,17 @@ import { GenkitError, Operation, z, - type Genkit, } from 'genkit'; import { BackgroundModelAction, - modelRef, + modelRef as createModelRef, type GenerateRequest, type ModelInfo, type ModelReference, } from 'genkit/model'; import { getApiKeyFromEnvVar } from './common.js'; import { Operation as ApiOperation, checkOp, predictModel } from './predict.js'; +import { backgroundModel } from 'genkit/plugin'; export type KNOWN_VEO_MODELS = 'veo-2.0-generate-001'; @@ -127,7 +127,6 @@ export const GENERIC_VEO_INFO = { } as ModelInfo; export function defineVeoModel( - ai: Genkit, name: string, apiKey?: string | false ): BackgroundModelAction { @@ -143,7 +142,7 @@ export function defineVeoModel( } } const modelName = `googleai/${name}`; - const model: ModelReference = modelRef({ + const modelRef: ModelReference = createModelRef({ name: modelName, info: { ...GENERIC_VEO_INFO, @@ -152,9 +151,9 @@ export function defineVeoModel( configSchema: VeoConfigSchema, }); - return ai.defineBackgroundModel({ + return backgroundModel({ name: modelName, - ...model.info, + ...modelRef.info, configSchema: VeoConfigSchema, async start(request) { const instance: VeoInstance = { @@ -169,7 +168,7 @@ export function defineVeoModel( VeoInstance, ApiOperation, VeoParameters - >(model.version || name, apiKey as string, 'predictLongRunning'); + >(modelRef.version || name, apiKey as string, 'predictLongRunning'); const response = await predictClient([instance], toParameters(request)); return toGenkitOp(response); From c79c78bee8da192eb5e6b344b81b49ef935179f2 Mon Sep 17 00:00:00 2001 From: Corie Watson Date: Mon, 27 Oct 2025 11:30:20 +0000 Subject: [PATCH 2/3] chore: format --- js/plugins/googleai/src/embedder.ts | 2 +- js/plugins/googleai/src/gemini.ts | 20 ++++---- js/plugins/googleai/src/imagen.ts | 4 +- js/plugins/googleai/src/index.ts | 71 +++++++++++++++++------------ js/plugins/googleai/src/veo.ts | 9 +--- 5 files changed, 59 insertions(+), 47 deletions(-) diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index ea90c51c40..0bb36b9455 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -25,9 +25,9 @@ import { type EmbedderReference, } from 'genkit'; import { embedderRef as createEmbedderRef } from 'genkit/embedder'; +import { embedder } from 'genkit/plugin'; import { getApiKeyFromEnvVar } from './common.js'; import type { PluginOptions } from './index.js'; -import { embedder } from 'genkit/plugin'; export const TaskTypeSchema = z.enum([ 'RETRIEVAL_DOCUMENT', diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index cb81227078..f7f7b933b7 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -42,8 +42,8 @@ import { GenkitError, z, type JSONSchema } from 'genkit'; import { GenerationCommonConfigDescriptions, GenerationCommonConfigSchema, - getBasicUsageStats, modelRef as createModelRef, + getBasicUsageStats, type CandidateData, type MediaPart, type MessageData, @@ -57,11 +57,11 @@ import { type ToolResponsePart, } from 'genkit/model'; import { downloadRequestMedia } from 'genkit/model/middleware'; +import { model } from 'genkit/plugin'; import { runInNewSpan } from 'genkit/tracing'; import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common'; import { handleCacheIfNeeded } from './context-caching'; import { extractCacheConfig } from './context-caching/utils'; -import { model } from 'genkit/plugin'; // Extra type guard to keep the compiler happy and avoid a cast to any. The // legacy Gemini SDK is no longer maintained, and doesn't have updated types. @@ -1205,7 +1205,10 @@ export function defineGoogleAIModel({ configSchema: GeminiConfigSchema, use: middleware, }, - async (request, { streamingRequested, sendChunk, abortSignal, registry }) => { + async ( + request, + { streamingRequested, sendChunk, abortSignal, registry } + ) => { const options: RequestOptions = { apiClient: getGenkitClientHeader() }; if (apiVersion) { options.apiVersion = apiVersion; @@ -1254,12 +1257,11 @@ export function defineGoogleAIModel({ if (codeExecutionFromConfig) { tools.push({ - codeExecution: - requestConfig.codeExecution ? - requestConfig.codeExecution === true - ? {} - : requestConfig.codeExecution - : {}, + codeExecution: requestConfig.codeExecution + ? requestConfig.codeExecution === true + ? {} + : requestConfig.codeExecution + : {}, }); } diff --git a/js/plugins/googleai/src/imagen.ts b/js/plugins/googleai/src/imagen.ts index 5209f42be3..3026776930 100644 --- a/js/plugins/googleai/src/imagen.ts +++ b/js/plugins/googleai/src/imagen.ts @@ -16,16 +16,16 @@ import { GenkitError, MessageData, z } from 'genkit'; import { - getBasicUsageStats, modelRef as createModelRef, + getBasicUsageStats, type GenerateRequest, type ModelAction, type ModelInfo, type ModelReference, } from 'genkit/model'; +import { model } from 'genkit/plugin'; import { getApiKeyFromEnvVar } from './common.js'; import { predictModel } from './predict.js'; -import { model } from 'genkit/plugin'; export type KNOWN_IMAGEN_MODELS = 'imagen-3.0-generate-002'; diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts index 092c6a81c1..59c8711cd1 100644 --- a/js/plugins/googleai/src/index.ts +++ b/js/plugins/googleai/src/index.ts @@ -25,7 +25,11 @@ import { } from 'genkit'; import { logger } from 'genkit/logging'; import { modelRef } from 'genkit/model'; -import { genkitPluginV2, ResolvableAction, type GenkitPluginV2 } from 'genkit/plugin'; +import { + ResolvableAction, + genkitPluginV2, + type GenkitPluginV2, +} from 'genkit/plugin'; import type { ActionType } from 'genkit/registry'; import { getApiKeyFromEnvVar } from './common.js'; import { @@ -111,7 +115,9 @@ export interface PluginOptions { experimental_debugTraces?: boolean; } -async function initializer(options?: PluginOptions): Promise { +async function initializer( + options?: PluginOptions +): Promise { let apiVersions = ['v1']; const actions: ResolvableAction[] = []; @@ -125,24 +131,28 @@ async function initializer(options?: PluginOptions): Promise 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, - })) + 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, - apiVersion: undefined, - baseUrl: options?.baseUrl, - debugTraces: options?.experimental_debugTraces, - })) + actions.push( + defineGoogleAIModel({ + name, + apiKey: options?.apiKey, + apiVersion: undefined, + baseUrl: options?.baseUrl, + debugTraces: options?.experimental_debugTraces, + }) + ) ); Object.keys(EMBEDDER_MODELS).forEach((name) => actions.push(defineGoogleAIEmbedder(name, { apiKey: options?.apiKey })) @@ -158,16 +168,18 @@ async function initializer(options?: PluginOptions): Promise modelOrRef.name.split('/')[1]; const modelRef = typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef; - actions.push(defineGoogleAIModel({ - name: modelName, - apiKey: options?.apiKey, - baseUrl: options?.baseUrl, - info: { - ...modelRef.info, - label: `Google AI - ${modelName}`, - }, - debugTraces: options?.experimental_debugTraces, - })); + actions.push( + defineGoogleAIModel({ + name: modelName, + apiKey: options?.apiKey, + baseUrl: options?.baseUrl, + info: { + ...modelRef.info, + label: `Google AI - ${modelName}`, + }, + debugTraces: options?.experimental_debugTraces, + }) + ); } } @@ -195,7 +207,10 @@ async function resolver( return undefined; } -function resolveModel(actionName: string, options?: PluginOptions): ResolvableAction | undefined { +function resolveModel( + actionName: string, + options?: PluginOptions +): ResolvableAction | undefined { if (actionName.startsWith('imagen')) { return defineImagenModel(actionName, options?.apiKey); } @@ -336,7 +351,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 { if (listActionsCache) return listActionsCache; listActionsCache = await listActions(options); return listActionsCache; - } + }, }); } diff --git a/js/plugins/googleai/src/veo.ts b/js/plugins/googleai/src/veo.ts index 9049701a37..9fb9aa07b0 100644 --- a/js/plugins/googleai/src/veo.ts +++ b/js/plugins/googleai/src/veo.ts @@ -14,12 +14,7 @@ * limitations under the License. */ -import { - GenerateResponseData, - GenkitError, - Operation, - z, -} from 'genkit'; +import { GenerateResponseData, GenkitError, Operation, z } from 'genkit'; import { BackgroundModelAction, modelRef as createModelRef, @@ -27,9 +22,9 @@ 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'; -import { backgroundModel } from 'genkit/plugin'; export type KNOWN_VEO_MODELS = 'veo-2.0-generate-001'; From ea103d710acf6eed820f23d3c64ab22ed7fe29bd Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Tue, 28 Oct 2025 15:04:47 +0000 Subject: [PATCH 3/3] refactor: update runInNewSpan call and tidy up init code --- js/plugins/googleai/src/gemini.ts | 8 ++------ js/plugins/googleai/src/index.ts | 14 ++++++++------ 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index f7f7b933b7..487de8bafb 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/js/plugins/googleai/src/gemini.ts @@ -1205,10 +1205,7 @@ export function defineGoogleAIModel({ configSchema: GeminiConfigSchema, use: middleware, }, - async ( - request, - { streamingRequested, sendChunk, abortSignal, registry } - ) => { + async (request, { streamingRequested, sendChunk, abortSignal }) => { const options: RequestOptions = { apiClient: getGenkitClientHeader() }; if (apiVersion) { options.apiVersion = apiVersion; @@ -1430,9 +1427,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 && registry + return debugTraces ? await runInNewSpan( - registry, { metadata: { name: streamingRequested ? 'sendMessageStream' : 'sendMessage', diff --git a/js/plugins/googleai/src/index.ts b/js/plugins/googleai/src/index.ts index 59c8711cd1..219fbae62b 100644 --- a/js/plugins/googleai/src/index.ts +++ b/js/plugins/googleai/src/index.ts @@ -130,8 +130,8 @@ async function initializer( } if (apiVersions.includes('v1beta')) { - Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => - actions.push( + actions.push( + ...Object.keys(SUPPORTED_GEMINI_MODELS).map((name) => defineGoogleAIModel({ name, apiKey: options?.apiKey, @@ -143,8 +143,8 @@ async function initializer( ); } if (apiVersions.includes('v1')) { - Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => - actions.push( + actions.push( + ...Object.keys(SUPPORTED_GEMINI_MODELS).map((name) => defineGoogleAIModel({ name, apiKey: options?.apiKey, @@ -154,8 +154,10 @@ async function initializer( }) ) ); - Object.keys(EMBEDDER_MODELS).forEach((name) => - actions.push(defineGoogleAIEmbedder(name, { apiKey: options?.apiKey })) + actions.push( + ...Object.keys(EMBEDDER_MODELS).map((name) => + defineGoogleAIEmbedder(name, { apiKey: options?.apiKey }) + ) ); }