Skip to content

Commit ff8ee3d

Browse files
ifielkerifielker
authored
feat(js/plugins/google-genai): Support for Nano Banana Pro (#3852)
1 parent c63f1c4 commit ff8ee3d

File tree

10 files changed

+375
-24
lines changed

10 files changed

+375
-24
lines changed

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

Lines changed: 81 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,31 @@ export const GeminiTtsConfigSchema = GeminiConfigSchema.extend({
314314
export type GeminiTtsConfigSchemaType = typeof GeminiTtsConfigSchema;
315315
export type GeminiTtsConfig = z.infer<GeminiTtsConfigSchemaType>;
316316

317+
export const GeminiImageConfigSchema = GeminiConfigSchema.extend({
318+
imageConfig: z
319+
.object({
320+
aspectRatio: z
321+
.enum([
322+
'1:1',
323+
'2:3',
324+
'3:2',
325+
'3:4',
326+
'4:3',
327+
'4:5',
328+
'5:4',
329+
'9:16',
330+
'16:9',
331+
'21:9',
332+
])
333+
.optional(),
334+
imageSize: z.enum(['1K', '2K', '4K']).optional(),
335+
})
336+
.passthrough()
337+
.optional(),
338+
}).passthrough();
339+
export type GeminiImageConfigSchemaType = typeof GeminiImageConfigSchema;
340+
export type GeminiImageConfig = z.infer<GeminiImageConfigSchemaType>;
341+
317342
export const GemmaConfigSchema = GeminiConfigSchema.extend({
318343
temperature: z
319344
.number()
@@ -332,7 +357,9 @@ export type GemmaConfig = z.infer<GemmaConfigSchemaType>;
332357
type ConfigSchemaType =
333358
| GeminiConfigSchemaType
334359
| GeminiTtsConfigSchemaType
360+
| GeminiImageConfigSchemaType
335361
| GemmaConfigSchemaType;
362+
type ConfigSchema = z.infer<ConfigSchemaType>;
336363

337364
function commonRef(
338365
name: string,
@@ -371,6 +398,20 @@ const GENERIC_TTS_MODEL = commonRef(
371398
},
372399
GeminiTtsConfigSchema
373400
);
401+
const GENERIC_IMAGE_MODEL = commonRef(
402+
'gemini-image',
403+
{
404+
supports: {
405+
multiturn: true,
406+
media: true,
407+
tools: true,
408+
toolChoice: true,
409+
systemRole: true,
410+
constrained: 'no-tools',
411+
},
412+
},
413+
GeminiImageConfigSchema
414+
);
374415
const GENERIC_GEMMA_MODEL = commonRef(
375416
'gemma-generic',
376417
undefined,
@@ -382,15 +423,17 @@ const KNOWN_GEMINI_MODELS = {
382423
'gemini-2.5-pro': commonRef('gemini-2.5-pro'),
383424
'gemini-2.5-flash': commonRef('gemini-2.5-flash'),
384425
'gemini-2.5-flash-lite': commonRef('gemini-2.5-flash-lite'),
385-
'gemini-2.5-flash-image-preview': commonRef('gemini-2.5-flash-image-preview'),
386-
'gemini-2.5-flash-image': commonRef('gemini-2.5-flash-image'),
387426
'gemini-2.0-flash': commonRef('gemini-2.0-flash'),
388427
'gemini-2.0-flash-lite': commonRef('gemini-2.0-flash-lite'),
389428
};
390429
export type KnownGeminiModels = keyof typeof KNOWN_GEMINI_MODELS;
391430
export type GeminiModelName = `gemini-${string}`;
392431
export function isGeminiModelName(value: string): value is GeminiModelName {
393-
return value.startsWith('gemini-') && !value.endsWith('-tts');
432+
return (
433+
value.startsWith('gemini-') &&
434+
!value.endsWith('-tts') &&
435+
!value.includes('-image')
436+
);
394437
}
395438

396439
const KNOWN_TTS_MODELS = {
@@ -411,6 +454,29 @@ export function isTTSModelName(value: string): value is TTSModelName {
411454
return value.startsWith('gemini-') && value.endsWith('-tts');
412455
}
413456

457+
const KNOWN_IMAGE_MODELS = {
458+
'gemini-3-pro-image-preview': commonRef(
459+
'gemini-3-pro-image-preview',
460+
{ ...GENERIC_IMAGE_MODEL.info },
461+
GeminiImageConfigSchema
462+
),
463+
'gemini-2.5-flash-image-preview': commonRef(
464+
'gemini-2.5-flash-image-preview',
465+
{ ...GENERIC_IMAGE_MODEL.info },
466+
GeminiImageConfigSchema
467+
),
468+
'gemini-2.5-flash-image': commonRef(
469+
'gemini-2.5-flash-image',
470+
{ ...GENERIC_IMAGE_MODEL.info },
471+
GeminiImageConfigSchema
472+
),
473+
} as const;
474+
export type KnownImageModels = keyof typeof KNOWN_IMAGE_MODELS;
475+
export type ImageModelName = `gemini-${string}-image${string}`;
476+
export function isImageModelName(value: string): value is ImageModelName {
477+
return value.startsWith('gemini-') && value.includes('-image');
478+
}
479+
414480
const KNOWN_GEMMA_MODELS = {
415481
'gemma-3-12b-it': commonRef('gemma-3-12b-it', undefined, GemmaConfigSchema),
416482
'gemma-3-1b-it': commonRef('gemma-3-1b-it', undefined, GemmaConfigSchema),
@@ -427,12 +493,13 @@ export function isGemmaModelName(value: string): value is GemmaModelName {
427493
const KNOWN_MODELS = {
428494
...KNOWN_GEMINI_MODELS,
429495
...KNOWN_TTS_MODELS,
496+
...KNOWN_IMAGE_MODELS,
430497
...KNOWN_GEMMA_MODELS,
431498
};
432499

433500
export function model(
434501
version: string,
435-
config: GeminiConfig | GeminiTtsConfig | GemmaConfig = {}
502+
config: ConfigSchema = {}
436503
): ModelReference<ConfigSchemaType> {
437504
const name = checkModelName(version);
438505

@@ -445,6 +512,15 @@ export function model(
445512
});
446513
}
447514

515+
if (isImageModelName(name)) {
516+
return modelRef({
517+
name: `googleai/${name}`,
518+
config,
519+
configSchema: GeminiImageConfigSchema,
520+
info: { ...GENERIC_IMAGE_MODEL.info },
521+
});
522+
}
523+
448524
if (isGemmaModelName(name)) {
449525
return modelRef({
450526
name: `googleai/${name}`,
@@ -562,7 +638,7 @@ export function defineModel(
562638
});
563639
}
564640

565-
const requestOptions: z.infer<ConfigSchemaType> = {
641+
const requestOptions: ConfigSchema = {
566642
...request.config,
567643
};
568644
const {

‎js/plugins/google-genai/src/googleai/index.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ async function resolver(
6161
} else if (imagen.isImagenModelName(actionName)) {
6262
return await imagen.defineModel(actionName, options);
6363
} else {
64-
// gemini, tts, gemma, unknown models
64+
// gemini, tts, image, gemma, unknown models
6565
return await gemini.defineModel(actionName, options);
6666
}
6767
break;
@@ -129,6 +129,10 @@ export type GoogleAIPlugin = {
129129
name: gemini.KnownTtsModels | (gemini.TTSModelName & {}),
130130
config: gemini.GeminiTtsConfig
131131
): ModelReference<gemini.GeminiTtsConfigSchemaType>;
132+
model(
133+
name: gemini.KnownImageModels | (gemini.ImageModelName & {}),
134+
config: gemini.GeminiImageConfig
135+
): ModelReference<gemini.GeminiImageConfigSchemaType>;
132136
model(
133137
name: gemini.KnownGeminiModels | (gemini.GeminiModelName & {}),
134138
config?: gemini.GeminiConfig
@@ -163,7 +167,7 @@ export const googleAI = googleAIPlugin as GoogleAIPlugin;
163167
if (imagen.isImagenModelName(name)) {
164168
return imagen.model(name, config);
165169
}
166-
// gemma, tts, gemini and unknown model families.
170+
// gemma, tts, image, gemini and unknown model families.
167171
return gemini.model(name, config);
168172
};
169173
googleAI.embedder = (

‎js/plugins/google-genai/src/vertexai/gemini.ts‎

Lines changed: 84 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -364,8 +364,34 @@ export type GeminiConfigSchemaType = typeof GeminiConfigSchema;
364364
*/
365365
export type GeminiConfig = z.infer<GeminiConfigSchemaType>;
366366

367+
export const GeminiImageConfigSchema = GeminiConfigSchema.extend({
368+
imageConfig: z
369+
.object({
370+
aspectRatio: z
371+
.enum([
372+
'1:1',
373+
'2:3',
374+
'3:2',
375+
'3:4',
376+
'4:3',
377+
'4:5',
378+
'5:4',
379+
'9:16',
380+
'16:9',
381+
'21:9',
382+
])
383+
.optional(),
384+
imageSize: z.enum(['1K', '2K', '4K']).optional(),
385+
})
386+
.passthrough()
387+
.optional(),
388+
}).passthrough();
389+
export type GeminiImageConfigSchemaType = typeof GeminiImageConfigSchema;
390+
export type GeminiImageConfig = z.infer<GeminiImageConfigSchemaType>;
391+
367392
// This contains all the Gemini config schema types
368-
type ConfigSchemaType = GeminiConfigSchemaType;
393+
type ConfigSchemaType = GeminiConfigSchemaType | GeminiImageConfigSchemaType;
394+
type ConfigSchema = z.infer<ConfigSchemaType>;
369395

370396
function commonRef(
371397
name: string,
@@ -388,9 +414,14 @@ function commonRef(
388414
});
389415
}
390416

391-
export const GENERIC_MODEL = commonRef('gemini');
417+
const GENERIC_MODEL = commonRef('gemini');
418+
const GENERIC_IMAGE_MODEL = commonRef(
419+
'gemini-image',
420+
undefined,
421+
GeminiImageConfigSchema
422+
);
392423

393-
export const KNOWN_MODELS = {
424+
export const KNOWN_GEMINI_MODELS = {
394425
'gemini-3-pro-preview': commonRef('gemini-3-pro-preview'),
395426
'gemini-2.5-flash-lite': commonRef('gemini-2.5-flash-lite'),
396427
'gemini-2.5-pro': commonRef('gemini-2.5-pro'),
@@ -400,21 +431,58 @@ export const KNOWN_MODELS = {
400431
'gemini-2.0-flash-lite': commonRef('gemini-2.0-flash-lite'),
401432
'gemini-2.0-flash-lite-001': commonRef('gemini-2.0-flash-lite-001'),
402433
} as const;
403-
export type KnownModels = keyof typeof KNOWN_MODELS;
434+
export type KnownGeminiModels = keyof typeof KNOWN_GEMINI_MODELS;
404435
export type GeminiModelName = `gemini-${string}`;
405436
export function isGeminiModelName(value?: string): value is GeminiModelName {
406-
return !!value?.startsWith('gemini-') && !value.includes('embedding');
437+
return !!(
438+
value?.startsWith('gemini-') &&
439+
!value.includes('embedding') &&
440+
!value.includes('-image')
441+
);
407442
}
408443

444+
export const KNOWN_IMAGE_MODELS = {
445+
'gemini-3-pro-image-preview': commonRef(
446+
'gemini-3-pro-image-preview',
447+
{ ...GENERIC_IMAGE_MODEL.info },
448+
GeminiImageConfigSchema
449+
),
450+
'gemini-2.5-flash-image': commonRef(
451+
'gemini-2.5-flash-image',
452+
undefined,
453+
GeminiImageConfigSchema
454+
),
455+
} as const;
456+
export type KnownImageModels = keyof typeof KNOWN_IMAGE_MODELS;
457+
export type ImageModelName = `gemini-${string}-image${string}`;
458+
export function isImageModelName(value?: string): value is ImageModelName {
459+
return !!(value?.startsWith('gemini-') && value.includes('-image'));
460+
}
461+
462+
const KNOWN_MODELS = {
463+
...KNOWN_GEMINI_MODELS,
464+
...KNOWN_IMAGE_MODELS,
465+
};
466+
export type KnownModels = keyof typeof KNOWN_MODELS;
467+
409468
export function model(
410469
version: string,
411-
options: GeminiConfig = {}
412-
): ModelReference<typeof GeminiConfigSchema> {
470+
config: ConfigSchema = {}
471+
): ModelReference<ConfigSchemaType> {
413472
const name = checkModelName(version);
414473

474+
if (isImageModelName(name)) {
475+
return modelRef({
476+
name: `vertexai/${name}`,
477+
config,
478+
configSchema: GeminiImageConfigSchema,
479+
info: { ...GENERIC_IMAGE_MODEL.info },
480+
});
481+
}
482+
415483
return modelRef({
416484
name: `vertexai/${name}`,
417-
config: options,
485+
config,
418486
configSchema: GeminiConfigSchema,
419487
info: {
420488
...GENERIC_MODEL.info,
@@ -433,7 +501,8 @@ export function listActions(models: Model[]): ActionMetadata[] {
433501
return models
434502
.filter(
435503
(m) =>
436-
isGeminiModelName(modelName(m.name)) &&
504+
(isGeminiModelName(modelName(m.name)) ||
505+
isImageModelName(modelName(m.name))) &&
437506
!KNOWN_DECOMISSIONED_MODELS.includes(modelName(m.name) || '')
438507
)
439508
.map((m) => {
@@ -509,7 +578,7 @@ export function defineModel(
509578
systemInstruction = toGeminiSystemInstruction(systemMessage);
510579
}
511580

512-
const requestConfig = { ...request.config };
581+
const requestConfig: ConfigSchema = { ...request.config };
513582

514583
const {
515584
apiKey: apiKeyFromConfig,
@@ -735,4 +804,8 @@ export function defineModel(
735804
);
736805
}
737806

738-
export const TEST_ONLY = { KNOWN_MODELS };
807+
export const TEST_ONLY = {
808+
KNOWN_GEMINI_MODELS,
809+
KNOWN_IMAGE_MODELS,
810+
KNOWN_MODELS,
811+
};

‎js/plugins/google-genai/src/vertexai/index.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,11 @@ function vertexAIPlugin(options?: VertexPluginOptions): GenkitPluginV2 {
124124
export type VertexAIPlugin = {
125125
(pluginOptions?: VertexPluginOptions): GenkitPluginV2;
126126
model(
127-
name: gemini.KnownModels | (gemini.GeminiModelName & {}),
127+
name: gemini.KnownImageModels | (gemini.ImageModelName & {}),
128+
config?: gemini.GeminiImageConfig
129+
): ModelReference<gemini.GeminiImageConfigSchemaType>;
130+
model(
131+
name: gemini.KnownGeminiModels | (gemini.GeminiModelName & {}),
128132
config?: gemini.GeminiConfig
129133
): ModelReference<gemini.GeminiConfigSchemaType>;
130134
model(
@@ -166,7 +170,7 @@ export const vertexAI = vertexAIPlugin as VertexAIPlugin;
166170
if (veo.isVeoModelName(name)) {
167171
return veo.model(name, config);
168172
}
169-
// gemini and unknown model families
173+
// gemini, image and unknown model families
170174
return gemini.model(name, config);
171175
};
172176
vertexAI.embedder = (

0 commit comments

Comments
 (0)