feat(js/plugins/google-genai): Added Veo models (#3193)

Author: ifielker

js/plugins/google-genai/src/common/utils.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { JSONSchema, ModelReference } from 'genkit';
+import { GenkitError, JSONSchema, ModelReference } from 'genkit';
 import { GenerateRequest, GenerationCommonConfigSchema } from 'genkit/model';
 import { ImagenInstance } from './types';
 
@@ -50,6 +50,23 @@ export function modelName(name?: string): string | undefined {
   return name?.split('/').at(-1);
 }
 
+/**
+ * Gets the suffix of a model string.
+ * Throws if the string is empty.
+ * @param name A string containing the model string
+ * @returns the model string stripped of prefixes and guaranteed not empty.
+ */
+export function checkModelName(name?: string): string {
+  const version = modelName(name);
+  if (!version) {
+    throw new GenkitError({
+      status: 'INVALID_ARGUMENT',
+      message: 'Model name is required.',
+    });
+  }
+  return version;
+}
+
 export function extractText(request: GenerateRequest) {
   return request.messages
     .at(-1)!

js/plugins/google-genai/src/googleai/client.ts

@@ -29,6 +29,8 @@ import {
   ListModelsResponse,
   Model,
   Part,
+  VeoOperation,
+  VeoPredictRequest,
 } from './types';
 
 /**
@@ -174,6 +176,46 @@ export async function imagenPredict(
   return response.json() as Promise<ImagenPredictResponse>;
 }
 
+export async function veoPredict(
+  apiKey: string,
+  model: string,
+  veoPredictRequest: VeoPredictRequest,
+  clientOptions?: ClientOptions
+): Promise<VeoOperation> {
+  const url = getGoogleAIUrl({
+    resourcePath: `models/${model}`,
+    resourceMethod: 'predictLongRunning',
+    clientOptions,
+  });
+
+  const fetchOptions: RequestInit = {
+    method: 'POST',
+    headers: await getHeaders(apiKey, clientOptions),
+    body: JSON.stringify(veoPredictRequest),
+  };
+
+  const response = await makeRequest(url, fetchOptions);
+  return response.json() as Promise<VeoOperation>;
+}
+
+export async function checkVeoOperation(
+  apiKey: string,
+  operation: string,
+  clientOptions?: ClientOptions
+): Promise<VeoOperation> {
+  const url = getGoogleAIUrl({
+    resourcePath: operation,
+    clientOptions,
+  });
+  const fetchOptions: RequestInit = {
+    method: 'GET',
+    headers: await getHeaders(apiKey, clientOptions),
+  };
+
+  const response = await makeRequest(url, fetchOptions);
+  return response.json() as Promise<VeoOperation>;
+}
+
 /**
  * Generates a Google AI URL.
  *

js/plugins/google-genai/src/googleai/embedder.ts

@@ -21,7 +21,6 @@ import {
   EmbedderInfo,
   EmbedderReference,
   Genkit,
-  GenkitError,
   z,
 } from 'genkit';
 import { embedderRef } from 'genkit/embedder';
@@ -32,7 +31,7 @@ import {
   Model,
   TaskTypeSchema,
 } from './types';
-import { calculateApiKey, checkApiKey, modelName } from './utils';
+import { calculateApiKey, checkApiKey, checkModelName } from './utils';
 
 export const EmbeddingConfigSchema = z
   .object({
@@ -80,26 +79,17 @@ function commonRef(
 
 const GENERIC_MODEL = commonRef('embedder');
 
-// TODO(ifielker): Update embedders to be current models.
-// (textEmbeddingGecko001 is outdated).
-
 const KNOWN_MODELS = {
   'text-embedding-004': commonRef('text-embedding-004'),
+  'gemini-embedding-exp': commonRef('gemini-embedding-exp'),
 };
 export type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete
 
 export function model(
   version: string,
   config: EmbeddingConfig = {}
 ): EmbedderReference<ConfigSchemaType> {
-  const name = modelName(version);
-  if (!name) {
-    throw new GenkitError({
-      status: 'INVALID_ARGUMENT',
-      message: 'Not able to create embedderReference for empty model version',
-    });
-  }
-
+  const name = checkModelName(version);
   return embedderRef({
     name: `googleai/${name}`,
     version: name,

js/plugins/google-genai/src/googleai/gemini.ts

@@ -59,7 +59,7 @@ import {
   Tool,
   ToolConfig,
 } from './types';
-import { calculateApiKey, checkApiKey, modelName } from './utils';
+import { calculateApiKey, checkApiKey, checkModelName } from './utils';
 
 /**
  * See https://ai.google.dev/gemini-api/docs/safety-settings#safety-filters.
@@ -298,14 +298,28 @@ const GENERIC_GEMMA_MODEL = commonRef(
   GemmaConfigSchema
 );
 
-const KNOWN_MODELS = {
+const KNOWN_GEMINI_MODELS = {
+  'gemini-2.5-pro': commonRef('gemini-2.5-pro'),
+  'gemini-2.5-flash': commonRef('gemini-2.5-flash'),
+  'gemini-2.5-flash-lite-preview-06-17': commonRef(
+    'gemini-2.5-flash-lite-preview-06-17'
+  ),
   'gemini-2.0-flash': commonRef('gemini-2.0-flash'),
+  'gemini-2.0-flash-preview-image-generation': commonRef(
+    'gemini-2.0-flash-preview-image-generation'
+  ),
   'gemini-2.0-flash-lite': commonRef('gemini-2.0-flash-lite'),
-  'gemini-2.0-pro-exp-02-05': commonRef('gemini-2.0-pro-exp-02-05'),
-  'gemini-2.0-flash-exp': commonRef('gemini-2.0-flash-exp'),
-  'gemini-2.5-pro-exp-03-25': commonRef('gemini-2.5-pro-exp-03-25'),
-  'gemini-2.5-pro-preview-03-25': commonRef('gemini-2.5-pro-preview-03-25'),
-  'gemini-2.5-flash-preview-04-17': commonRef('gemini-2.5-flash-preview-04-17'),
+  'gemini-1.5-flash': commonRef('gemini-1.5-flash'),
+  'gemini-1.5-flash-8b': commonRef('gemini-1.5-flash-8b'),
+  'gemini-1.5-pro': commonRef('gemini-1.5-pro'),
+};
+export type KnownGeminiModels = keyof typeof KNOWN_GEMINI_MODELS;
+export type GeminiModelName = `gemini-${string}`;
+export function isGeminiModelName(value: string): value is GeminiModelName {
+  return value.startsWith('gemini-') && !value.endsWith('-tts');
+}
+
+const KNOWN_TTS_MODELS = {
   'gemini-2.5-flash-preview-tts': commonRef(
     'gemini-2.5-flash-preview-tts',
     { ...GENERIC_TTS_MODEL.info },
@@ -316,36 +330,37 @@ const KNOWN_MODELS = {
     { ...GENERIC_TTS_MODEL.info },
     GeminiTtsConfigSchema
   ),
+};
+export type KnownTtsModels = keyof typeof KNOWN_TTS_MODELS;
+export type TTSModelName = `gemini-${string}-tts`;
+export function isTTSModelName(value: string): value is TTSModelName {
+  return value.startsWith('gemini-') && value.endsWith('-tts');
+}
+
+const KNOWN_GEMMA_MODELS = {
   'gemma-3-12b-it': commonRef('gemma-3-12b-it', undefined, GemmaConfigSchema),
   'gemma-3-1b-it': commonRef('gemma-3-1b-it', undefined, GemmaConfigSchema),
   'gemma-3-27b-it': commonRef('gemma-3-27b-it', undefined, GemmaConfigSchema),
   'gemma-3-4b-it': commonRef('gemma-3-4b-it', undefined, GemmaConfigSchema),
   'gemma-3n-e4b-it': commonRef('gemma-3n-e4b-it', undefined, GemmaConfigSchema),
 } as const;
-export type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete
-
-// For conditional types in index.ts model()
-export type TTSModelName = `gemini-${string}-tts`;
-export function isTTSModelName(value: string): value is TTSModelName {
-  return value.startsWith('gemini-') && value.endsWith('-tts');
-}
-
+export type KnownGemmaModels = keyof typeof KNOWN_GEMMA_MODELS;
 export type GemmaModelName = `gemma-${string}`;
 export function isGemmaModelName(value: string): value is GemmaModelName {
   return value.startsWith('gemma-');
 }
 
+const KNOWN_MODELS = {
+  ...KNOWN_GEMINI_MODELS,
+  ...KNOWN_TTS_MODELS,
+  ...KNOWN_GEMMA_MODELS,
+};
+
 export function model(
   version: string,
   config: GeminiConfig | GeminiTtsConfig = {}
 ): ModelReference<ConfigSchemaType> {
-  const name = modelName(version);
-  if (!name) {
-    throw new GenkitError({
-      status: 'INVALID_ARGUMENT',
-      message: 'Not able to create modelReference for empty model version',
-    });
-  }
+  const name = checkModelName(version);
 
   if (isTTSModelName(name)) {
     return modelRef({
@@ -411,6 +426,10 @@ export function defineModel(
 ): ModelAction {
   checkApiKey(pluginOptions?.apiKey);
   const ref = model(name);
+  const clientOptions: ClientOptions = {
+    apiVersion: pluginOptions?.apiVersion,
+    baseUrl: pluginOptions?.baseUrl,
+  };
 
   const middleware: ModelMiddleware[] = [];
   if (ref.info?.supports?.media) {
@@ -447,11 +466,6 @@ export function defineModel(
       use: middleware,
     },
     async (request, sendChunk) => {
-      const clientOptions: ClientOptions = {
-        apiVersion: pluginOptions?.apiVersion,
-        baseUrl: pluginOptions?.baseUrl,
-      };
-
       // Make a copy so that modifying the request will not produce side-effects
       const messages = [...request.messages];
       if (messages.length === 0) throw new Error('No messages provided.');

js/plugins/google-genai/src/googleai/imagen.ts

@@ -16,7 +16,6 @@
 
 import {
   ActionMetadata,
-  GenkitError,
   MediaPart,
   MessageData,
   modelActionMetadata,
@@ -43,6 +42,7 @@ import type {
 import {
   calculateApiKey,
   checkApiKey,
+  checkModelName,
   extractImagenImage,
   extractText,
   modelName,
@@ -116,6 +116,12 @@ const GENERIC_MODEL = commonRef('imagen', {
 
 const KNOWN_MODELS = {
   'imagen-3.0-generate-002': commonRef('imagen-3.0-generate-002'),
+  'imagen-4.0-generate-preview-06-06': commonRef(
+    'imagen-4.0-generate-preview-06-06'
+  ),
+  'imagen-4.0-ultra-generate-preview-06-06': commonRef(
+    'imagen-4.0-ultra-generate-preview-06-06'
+  ),
 } as const;
 export type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete
 
@@ -129,13 +135,7 @@ export function model(
   version: string,
   config: ImagenConfig = {}
 ): ModelReference<ConfigSchemaType> {
-  const name = modelName(version);
-  if (!name) {
-    throw new GenkitError({
-      status: 'INVALID_ARGUMENT',
-      message: 'Not able to create modelReference for empty model version',
-    });
-  }
+  const name = checkModelName(version);
   if (KNOWN_MODELS[name]) {
     return KNOWN_MODELS[name].withConfig(config);
   }