Skip to content

Commit f365d79

Browse files
ifielkerifielker
authored
fix(js/plugins/google-genai): Fixed imagen, veo, and streaming thoughts (#3279)
1 parent eb7729b commit f365d79

File tree

14 files changed

+78
-36
lines changed

14 files changed

+78
-36
lines changed

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

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

17-
import { GenkitError, JSONSchema } from 'genkit';
17+
import {
18+
EmbedderReference,
19+
GenkitError,
20+
JSONSchema,
21+
ModelReference,
22+
z,
23+
} from 'genkit';
1824
import { GenerateRequest } from 'genkit/model';
1925
import { ImagenInstance } from './types';
2026

@@ -41,13 +47,26 @@ export function extractErrMsg(e: unknown): string {
4147
}
4248

4349
/**
44-
* Gets the suffix of a model string.
50+
* Gets the un-prefixed model name from a modelReference
51+
*/
52+
export function extractVersion(
53+
model: ModelReference<z.ZodTypeAny> | EmbedderReference<z.ZodTypeAny>
54+
): string {
55+
return model.version ? model.version : checkModelName(model.name);
56+
}
57+
58+
/**
59+
* Gets the model name without certain prefixes..
4560
* e.g. for "models/googleai/gemini-2.5-pro" it returns just 'gemini-2.5-pro'
4661
* @param name A string containing the model string with possible prefixes
47-
* @returns the model string stripped of prefixes
62+
* @returns the model string stripped of certain prefixes
4863
*/
4964
export function modelName(name?: string): string | undefined {
50-
return name?.split('/').at(-1);
65+
if (!name) return name;
66+
67+
// Remove any of these prefixes: (but keep tunedModels e.g.)
68+
const prefixesToRemove = /models\/|embedders\/|googleai\/|vertexai\//g;
69+
return name.replace(prefixesToRemove, '');
5170
}
5271

5372
/**

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -444,6 +444,9 @@ function aggregateResponses(
444444

445445
for (const part of candidate.content.parts) {
446446
const newPart: Partial<Part> = {};
447+
if (part.thought) {
448+
newPart.thought = part.thought;
449+
}
447450
if (part.text) {
448451
newPart.text = part.text;
449452
}

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,12 @@ import {
3131
Model,
3232
TaskTypeSchema,
3333
} from './types';
34-
import { calculateApiKey, checkApiKey, checkModelName } from './utils';
34+
import {
35+
calculateApiKey,
36+
checkApiKey,
37+
checkModelName,
38+
extractVersion,
39+
} from './utils';
3540

3641
export const EmbeddingConfigSchema = z
3742
.object({
@@ -92,7 +97,6 @@ export function model(
9297
const name = checkModelName(version);
9398
return embedderRef({
9499
name: `googleai/${name}`,
95-
version: name,
96100
config,
97101
configSchema: GENERIC_MODEL.configSchema,
98102
info: {
@@ -143,7 +147,7 @@ export function defineEmbedder(
143147
pluginOptions?.apiKey,
144148
reqOptions?.apiKey
145149
);
146-
const embedVersion = (reqOptions?.version || ref.version) as string;
150+
const embedVersion = reqOptions?.version || extractVersion(ref);
147151

148152
const embeddings = await Promise.all(
149153
input.map(async (doc) => {

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

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@ import {
4040
toGeminiSystemInstruction,
4141
toGeminiTool,
4242
} from '../common/converters';
43-
import { cleanSchema } from '../common/utils';
4443
import {
4544
generateContent,
4645
generateContentStream,
@@ -59,7 +58,13 @@ import {
5958
Tool,
6059
ToolConfig,
6160
} from './types';
62-
import { calculateApiKey, checkApiKey, checkModelName } from './utils';
61+
import {
62+
calculateApiKey,
63+
checkApiKey,
64+
checkModelName,
65+
cleanSchema,
66+
extractVersion,
67+
} from './utils';
6368

6469
/**
6570
* See https://ai.google.dev/gemini-api/docs/safety-settings#safety-filters.
@@ -360,7 +365,6 @@ export function model(
360365
if (isTTSModelName(name)) {
361366
return modelRef({
362367
name: `googleai/${name}`,
363-
version: name,
364368
config,
365369
configSchema: GeminiTtsConfigSchema,
366370
info: { ...GENERIC_TTS_MODEL.info },
@@ -370,7 +374,6 @@ export function model(
370374
if (isGemmaModelName(name)) {
371375
return modelRef({
372376
name: `googleai/${name}`,
373-
version: name,
374377
config,
375378
configSchema: GemmaConfigSchema,
376379
info: { ...GENERIC_GEMMA_MODEL.info },
@@ -379,7 +382,6 @@ export function model(
379382

380383
return modelRef({
381384
name: `googleai/${name}`,
382-
version: name,
383385
config,
384386
configSchema: GeminiConfigSchema,
385387
info: { ...GENERIC_MODEL.info },
@@ -562,7 +564,7 @@ export function defineModel(
562564
contents: messages.map((message) => toGeminiMessage(message, ref)),
563565
};
564566

565-
const modelVersion = (versionFromConfig || ref.version) as string;
567+
const modelVersion = versionFromConfig || extractVersion(ref);
566568

567569
const generateApiKey = calculateApiKey(
568570
pluginOptions?.apiKey,

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import {
4545
checkModelName,
4646
extractImagenImage,
4747
extractText,
48+
extractVersion,
4849
modelName,
4950
} from './utils.js';
5051

@@ -142,7 +143,6 @@ export function model(
142143

143144
return modelRef({
144145
name: `googleai/${name}`,
145-
version: name,
146146
config,
147147
configSchema: ImagenConfigSchema,
148148
info: {
@@ -213,7 +213,7 @@ export function defineModel(
213213

214214
const response = await imagenPredict(
215215
predictApiKey,
216-
ref.version as string,
216+
extractVersion(ref),
217217
imagenPredictRequest,
218218
clientOpt
219219
);

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@ import { VeoImage } from './types.js';
2020

2121
export {
2222
checkModelName,
23+
cleanSchema,
2324
extractImagenImage,
2425
extractText,
26+
extractVersion,
2527
modelName,
2628
} from '../common/utils.js';
2729

‎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+
extractVersion,
4647
modelName,
4748
} from './utils.js';
4849

@@ -124,7 +125,6 @@ export function model(
124125
const name = checkModelName(version);
125126
return modelRef({
126127
name: `googleai/${name}`,
127-
version: name,
128128
config,
129129
configSchema: VeoConfigSchema,
130130
info: { ...GENERIC_MODEL.info },
@@ -192,7 +192,7 @@ export function defineModel(
192192

193193
const response = await veoPredict(
194194
apiKey,
195-
ref.version as string,
195+
extractVersion(ref),
196196
veoPredictRequest,
197197
clientOptions
198198
);
@@ -228,6 +228,10 @@ function toVeoParameters(
228228
// It's pulled out and used separately
229229
delete out.apiKey;
230230

231+
// This was used to help us figure out which model. We no longer need
232+
// it here.
233+
delete out.version;
234+
231235
return out;
232236
}
233237

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ import {
3333
isMultimodalEmbeddingPrediction,
3434
isObject,
3535
} from './types';
36-
import { checkModelName } from './utils';
36+
import { checkModelName, extractVersion } from './utils';
3737

3838
export const EmbeddingConfigSchema = z
3939
.object({
@@ -116,7 +116,6 @@ export function model(
116116
if (KNOWN_MODELS[name]) {
117117
return embedderRef({
118118
name: `vertexai/${name}`,
119-
version: name,
120119
configSchema: EmbeddingConfigSchema,
121120
config,
122121
info: {
@@ -128,7 +127,6 @@ export function model(
128127
// Generic multimodal embedder format
129128
return embedderRef({
130129
name: `vertexai/${name}`,
131-
version: name,
132130
configSchema: EmbeddingConfigSchema,
133131
config,
134132
info: {
@@ -139,7 +137,6 @@ export function model(
139137
// Generic text-only embedder format
140138
return embedderRef({
141139
name: `vertexai/${name}`,
142-
version: name,
143140
configSchema: EmbeddingConfigSchema,
144141
config,
145142
info: {
@@ -181,7 +178,7 @@ export function defineEmbedder(
181178
};
182179

183180
const response = await embedContent(
184-
ref.version as string,
181+
extractVersion(ref),
185182
embedContentRequest,
186183
clientOptions
187184
);

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

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@ import {
4040
toGeminiSystemInstruction,
4141
toGeminiTool,
4242
} from '../common/converters';
43-
import { checkModelName, cleanSchema, modelName } from '../common/utils';
4443
import {
4544
generateContent,
4645
generateContentStream,
@@ -59,7 +58,13 @@ import {
5958
ToolConfig,
6059
VertexPluginOptions,
6160
} from './types';
62-
import { calculateApiKey } from './utils';
61+
import {
62+
calculateApiKey,
63+
checkModelName,
64+
cleanSchema,
65+
extractVersion,
66+
modelName,
67+
} from './utils';
6368

6469
export const SafetySettingsSchema = z.object({
6570
category: z.enum([
@@ -301,7 +306,6 @@ function commonRef(
301306
): ModelReference<ConfigSchemaType> {
302307
return modelRef({
303308
name: `vertexai/${name}`,
304-
version: name,
305309
configSchema,
306310
info: info ?? {
307311
supports: {
@@ -341,7 +345,6 @@ export function model(
341345

342346
return modelRef({
343347
name: `vertexai/${name}`,
344-
version: name,
345348
config: options,
346349
configSchema: GeminiConfigSchema,
347350
info: {
@@ -583,7 +586,7 @@ export function defineModel(
583586
labels,
584587
};
585588

586-
const modelVersion = versionFromConfig || (ref.version as string);
589+
const modelVersion = versionFromConfig || extractVersion(ref);
587590

588591
if (jsonMode && request.output?.constrained) {
589592
generateContentRequest.generationConfig!.responseSchema = cleanSchema(

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import {
3939
extractImagenImage,
4040
extractImagenMask,
4141
extractText,
42+
extractVersion,
4243
modelName,
4344
} from './utils';
4445

@@ -238,7 +239,6 @@ export function model(
238239
}
239240
return modelRef({
240241
name: `vertexai/${name}`,
241-
version,
242242
config,
243243
configSchema: ImagenConfigSchema,
244244
info: {
@@ -299,7 +299,7 @@ export function defineModel(
299299
};
300300

301301
const response = await imagenPredict(
302-
ref.version as string,
302+
extractVersion(ref),
303303
imagenPredictRequest,
304304
clientOpt
305305
);

0 commit comments

Comments
 (0)