Skip to content

Commit a4f92b3

Browse files
authored
feat(js/plugins/google-genai): New models and support for googleai veo video extension (#3736)
1 parent 8e1c4c9 commit a4f92b3

File tree

13 files changed

+351
-43
lines changed

13 files changed

+351
-43
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,7 @@ const KNOWN_GEMINI_MODELS = {
333333
'gemini-2.5-flash': commonRef('gemini-2.5-flash'),
334334
'gemini-2.5-flash-lite': commonRef('gemini-2.5-flash-lite'),
335335
'gemini-2.5-flash-image-preview': commonRef('gemini-2.5-flash-image-preview'),
336+
'gemini-2.5-flash-image': commonRef('gemini-2.5-flash-image'),
336337
'gemini-2.0-flash': commonRef('gemini-2.0-flash'),
337338
'gemini-2.0-flash-preview-image-generation': commonRef(
338339
'gemini-2.0-flash-preview-image-generation'

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

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -117,12 +117,9 @@ const GENERIC_MODEL = commonRef('imagen', {
117117

118118
const KNOWN_MODELS = {
119119
'imagen-3.0-generate-002': commonRef('imagen-3.0-generate-002'),
120-
'imagen-4.0-generate-preview-06-06': commonRef(
121-
'imagen-4.0-generate-preview-06-06'
122-
),
123-
'imagen-4.0-ultra-generate-preview-06-06': commonRef(
124-
'imagen-4.0-ultra-generate-preview-06-06'
125-
),
120+
'imagen-4.0-fast-generate-001': commonRef('imagen-4.0-fast-generate-001'),
121+
'imagen-4.0-generate-001': commonRef('imagen-4.0-generate-001'),
122+
'imagen-4.0-ultra-generate-001': commonRef('imagen-4.0-ultra-generate-001'),
126123
} as const;
127124
export type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete
128125

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,13 +186,18 @@ export declare interface VeoParameters {
186186
export declare interface VeoInstance {
187187
prompt: string;
188188
image?: VeoImage;
189+
video?: VeoVideo;
189190
}
190191

191192
export declare interface VeoImage {
192193
bytesBase64Encoded: string;
193194
mimeType: string;
194195
}
195196

197+
export declare interface VeoVideo {
198+
uri: string;
199+
}
200+
196201
export declare interface VeoOperation {
197202
name: string;
198203
done?: boolean;

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

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import { GenerateRequest, GenkitError } from 'genkit';
1818
import process from 'process';
1919
import { extractMedia } from '../common/utils.js';
20-
import { ImagenInstance, VeoImage } from './types.js';
20+
import { ImagenInstance, VeoImage, VeoVideo } from './types.js';
2121

2222
export {
2323
checkModelName,
@@ -126,24 +126,31 @@ export function extractVeoImage(
126126
request: GenerateRequest
127127
): VeoImage | undefined {
128128
const media = request.messages.at(-1)?.content.find((p) => !!p.media)?.media;
129-
if (media) {
130-
const img = media.url.split(',')[1];
131-
if (img && media.contentType) {
132-
return {
133-
bytesBase64Encoded: img,
134-
mimeType: media.contentType!,
135-
};
136-
} else if (img) {
137-
// Content Type is not optional
138-
throw new GenkitError({
139-
status: 'INVALID_ARGUMENT',
140-
message: 'content type is required for images',
141-
});
142-
}
129+
if (!media?.contentType?.startsWith('image/')) {
130+
return undefined;
131+
}
132+
const bytes = media?.url.split(',')[1];
133+
if (bytes) {
134+
return {
135+
bytesBase64Encoded: bytes,
136+
mimeType: media.contentType,
137+
};
143138
}
144139
return undefined;
145140
}
146141

142+
export function extractVeoVideo(
143+
request: GenerateRequest
144+
): VeoVideo | undefined {
145+
const media = request.messages.at(-1)?.content.find((p) => !!p.media)?.media;
146+
if (!media?.contentType?.startsWith('video/')) {
147+
return undefined;
148+
}
149+
return {
150+
uri: media.url,
151+
};
152+
}
153+
147154
export function extractImagenImage(
148155
request: GenerateRequest
149156
): ImagenInstance['image'] | undefined {

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ import {
4343
checkModelName,
4444
extractText,
4545
extractVeoImage,
46+
extractVeoVideo,
4647
extractVersion,
4748
modelName,
4849
} from './utils.js';
@@ -110,8 +111,10 @@ function commonRef(
110111
const GENERIC_MODEL = commonRef('veo');
111112

112113
const KNOWN_MODELS = {
113-
'veo-3.0-generate-preview': commonRef('veo-3.0-generate-preview'),
114-
'veo-3.0-fast-generate-preview': commonRef('veo-3.0-fast-generate-preview'),
114+
'veo-3.1-generate-preview': commonRef('veo-3.1-generate-preview'),
115+
'veo-3.1-fast-generate-preview': commonRef('veo-3.1-fast-generate-preview'),
116+
'veo-3.0-generate-001': commonRef('veo-3.0-generate-001'),
117+
'veo-3.0-fast-generate-001': commonRef('veo-3.0-fast-generate-001'),
115118
'veo-2.0-generate-001': commonRef('veo-2.0-generate-001'),
116119
} as const;
117120
export type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete
@@ -186,6 +189,7 @@ export function defineModel(
186189
{
187190
prompt: extractText(request),
188191
image: extractVeoImage(request),
192+
video: extractVeoVideo(request),
189193
},
190194
],
191195
parameters: toVeoParameters(request),

js/plugins/google-genai/src/vertexai/veo.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,10 @@ const KNOWN_MODELS = {
136136
'veo-3.0-fast-generate-001': commonRef('veo-3.0-fast-generate-001'),
137137
'veo-3.0-generate-preview': commonRef('veo-3.0-generate-preview'),
138138
'veo-3.0-fast-generate-preview': commonRef('veo-3.0-fast-generate-preview'),
139+
'veo-3.1-fast-generate-001': commonRef('veo-3.1-fast-generate-001'),
140+
'veo-3.1-fast-generate-preview': commonRef('veo-3.1-fast-generate-preview'),
141+
'veo-3.1-generate-001': commonRef('veo-3.1-generate-001'),
142+
'veo-3.1-generate-preview': commonRef('veo-3.1-generate-preview'),
139143
} as const;
140144
export type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete
141145
export type VeoModelName = `veo-${string}`;

js/plugins/google-genai/tests/googleai/gemini_test.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,14 @@ describe('Google AI Gemini', () => {
346346
assert.strictEqual(modelRef.configSchema, GeminiTtsConfigSchema);
347347
});
348348

349+
it('returns a ModelReference for an image type model string', () => {
350+
const name = 'gemini-2.5-flash-image';
351+
const modelRef = model(name);
352+
assert.strictEqual(modelRef.name, `googleai/${name}`);
353+
assert.strictEqual(modelRef.info?.supports?.multiturn, true);
354+
assert.strictEqual(modelRef.configSchema, GeminiConfigSchema);
355+
});
356+
349357
it('returns a ModelReference for an unknown model string', () => {
350358
const name = 'gemini-3.0-flash';
351359
const modelRef = model(name);

js/plugins/google-genai/tests/googleai/imagen_test.ts

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ describe('Google AI Imagen', () => {
4949

5050
describe('model()', () => {
5151
it('should return a ModelReference for a known model', () => {
52-
const modelName = 'imagen-3.0-generate-002';
52+
const modelName = 'imagen-4.0-generate-001';
5353
const ref = model(modelName);
5454
assert.strictEqual(ref.name, `googleai/${modelName}`);
5555
assert.ok(ref.info?.supports?.media);
@@ -63,7 +63,7 @@ describe('Google AI Imagen', () => {
6363
});
6464

6565
it('should apply config to a known model', () => {
66-
const modelName = 'imagen-3.0-generate-002';
66+
const modelName = 'imagen-4.0-generate-001';
6767
const config: ImagenConfig = { numberOfImages: 2 };
6868
const ref = model(modelName, config);
6969
assert.strictEqual(ref.name, `googleai/${modelName}`);
@@ -79,9 +79,9 @@ describe('Google AI Imagen', () => {
7979
});
8080

8181
it('should handle model name with prefix', () => {
82-
const modelName = 'models/imagen-3.0-generate-002';
82+
const modelName = 'models/imagen-4.0-generate-001';
8383
const ref = model(modelName);
84-
assert.strictEqual(ref.name, 'googleai/imagen-3.0-generate-002');
84+
assert.strictEqual(ref.name, 'googleai/imagen-4.0-generate-001');
8585
});
8686
});
8787

js/plugins/google-genai/tests/googleai/index_test.ts

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -434,11 +434,11 @@ describe('GoogleAI Plugin', () => {
434434
supportedGenerationMethods: ['embedContent'],
435435
},
436436
{
437-
name: 'models/imagen-3.0-generate-002',
437+
name: 'models/imagen-4.0-generate-001',
438438
supportedGenerationMethods: ['predict'],
439439
},
440440
{
441-
name: 'models/veo-2.0-generate-001',
441+
name: 'models/veo-3.1-generate-preview',
442442
supportedGenerationMethods: ['predictLongRunning'],
443443
},
444444
{
@@ -457,9 +457,9 @@ describe('GoogleAI Plugin', () => {
457457
actionNames,
458458
[
459459
'googleai/gemini-2.5-pro',
460-
'googleai/imagen-3.0-generate-002',
460+
'googleai/imagen-4.0-generate-001',
461461
'googleai/text-embedding-004',
462-
'googleai/veo-2.0-generate-001',
462+
'googleai/veo-3.1-generate-preview',
463463
].sort()
464464
);
465465

@@ -474,12 +474,12 @@ describe('GoogleAI Plugin', () => {
474474
assert.strictEqual(embedderAction?.actionType, 'embedder');
475475

476476
const imagenAction = actions.find(
477-
(a) => a.name === 'googleai/imagen-3.0-generate-002'
477+
(a) => a.name === 'googleai/imagen-4.0-generate-001'
478478
);
479479
assert.strictEqual(imagenAction?.actionType, 'model');
480480

481481
const veoAction = actions.find(
482-
(a) => a.name === 'googleai/veo-2.0-generate-001'
482+
(a) => a.name === 'googleai/veo-3.1-generate-preview'
483483
);
484484
assert.strictEqual(veoAction?.actionType, 'model');
485485
});

js/plugins/google-genai/tests/googleai/utils_test.ts

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
* limitations under the License.
1515
*/
1616

17+
import { GenerateRequest } from 'genkit';
1718
import assert from 'node:assert';
1819
import { afterEach, beforeEach, describe, it } from 'node:test';
1920
import process from 'process';
@@ -22,9 +23,152 @@ import {
2223
MISSING_API_KEY_ERROR,
2324
calculateApiKey,
2425
checkApiKey,
26+
extractVeoImage,
27+
extractVeoVideo,
2528
getApiKeyFromEnvVar,
2629
} from '../../src/googleai/utils.js'; // Assuming the file is named utils.ts
2730

31+
describe('Media Utils', () => {
32+
describe('extractVeoImage', () => {
33+
it('should extract an image from a valid request', () => {
34+
const request: GenerateRequest = {
35+
messages: [
36+
{
37+
role: 'user',
38+
content: [
39+
{
40+
media: {
41+
url: 'data:image/jpeg;base64,test-base64-string',
42+
contentType: 'image/jpeg',
43+
},
44+
},
45+
],
46+
},
47+
],
48+
};
49+
const result = extractVeoImage(request);
50+
assert.deepStrictEqual(result, {
51+
bytesBase64Encoded: 'test-base64-string',
52+
mimeType: 'image/jpeg',
53+
});
54+
});
55+
56+
it('should return undefined if no media part is present', () => {
57+
const request: GenerateRequest = {
58+
messages: [
59+
{
60+
role: 'user',
61+
content: [{ text: 'no media here' }],
62+
},
63+
],
64+
};
65+
const result = extractVeoImage(request);
66+
assert.strictEqual(result, undefined);
67+
});
68+
69+
it('should return undefined if media is not an image', () => {
70+
const request: GenerateRequest = {
71+
messages: [
72+
{
73+
role: 'user',
74+
content: [
75+
{
76+
media: {
77+
url: 'data:video/mp4;base64,test-video',
78+
contentType: 'video/mp4',
79+
},
80+
},
81+
],
82+
},
83+
],
84+
};
85+
const result = extractVeoImage(request);
86+
assert.strictEqual(result, undefined);
87+
});
88+
89+
it('should extract from a malformed data URL', () => {
90+
const request: GenerateRequest = {
91+
messages: [
92+
{
93+
role: 'user',
94+
content: [
95+
{
96+
media: {
97+
url: 'data:image/jpeg,missing-base64-part', // Invalid format
98+
contentType: 'image/jpeg',
99+
},
100+
},
101+
],
102+
},
103+
],
104+
};
105+
const result = extractVeoImage(request);
106+
assert.deepStrictEqual(result, {
107+
bytesBase64Encoded: 'missing-base64-part',
108+
mimeType: 'image/jpeg',
109+
});
110+
});
111+
});
112+
113+
describe('extractVeoVideo', () => {
114+
it('should extract a video from a valid request', () => {
115+
const videoUrl = 'http://example.com/video.mp4';
116+
const request: GenerateRequest = {
117+
messages: [
118+
{
119+
role: 'user',
120+
content: [
121+
{
122+
media: {
123+
url: videoUrl,
124+
contentType: 'video/mp4',
125+
},
126+
},
127+
],
128+
},
129+
],
130+
};
131+
const result = extractVeoVideo(request);
132+
assert.deepStrictEqual(result, {
133+
uri: videoUrl,
134+
});
135+
});
136+
137+
it('should return undefined if no media part is present', () => {
138+
const request: GenerateRequest = {
139+
messages: [
140+
{
141+
role: 'user',
142+
content: [{ text: 'no media here' }],
143+
},
144+
],
145+
};
146+
const result = extractVeoVideo(request);
147+
assert.strictEqual(result, undefined);
148+
});
149+
150+
it('should return undefined if media is not a video', () => {
151+
const request: GenerateRequest = {
152+
messages: [
153+
{
154+
role: 'user',
155+
content: [
156+
{
157+
media: {
158+
url: 'data:image/jpeg;base64,test-image',
159+
contentType: 'image/jpeg',
160+
},
161+
},
162+
],
163+
},
164+
],
165+
};
166+
const result = extractVeoVideo(request);
167+
assert.strictEqual(result, undefined);
168+
});
169+
});
170+
});
171+
28172
describe('API Key Utils', () => {
29173
let originalEnv: NodeJS.ProcessEnv;
30174

0 commit comments

Comments
 (0)