Skip to content

Commit f66475b

Browse files
cabljaccabljac
committed
refactor(plugins/googleai): handle prefixes better
1 parent 18f0670 commit f66475b

File tree

7 files changed

+402
-29
lines changed

7 files changed

+402
-29
lines changed

‎js/plugins/googleai/src/common.ts

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,32 @@ export function getGenkitClientHeader() {
3131
}
3232
return defaultGetClientHeader();
3333
}
34+
35+
// Type-safe name helpers for GoogleAI plugin
36+
const PROVIDER = 'googleai' as const;
37+
type Provider = typeof PROVIDER;
38+
type Prefixed<ActionName extends string> = `${Provider}/${ActionName}`;
39+
type MaybePrefixed<ActionName extends string> =
40+
| ActionName
41+
| Prefixed<ActionName>;
42+
43+
// Runtime + typed helpers
44+
export function removePrefix<ActionName extends string>(
45+
name: MaybePrefixed<ActionName>
46+
): ActionName {
47+
return (
48+
name.startsWith(`${PROVIDER}/`)
49+
? (name.slice(PROVIDER.length + 1) as ActionName)
50+
: name
51+
) as ActionName;
52+
}
53+
54+
export function ensurePrefixed<ActionName extends string>(
55+
name: MaybePrefixed<ActionName>
56+
): Prefixed<ActionName> {
57+
return (
58+
name.startsWith(`${PROVIDER}/`)
59+
? (name as Prefixed<ActionName>)
60+
: (`${PROVIDER}/${name}` as Prefixed<ActionName>)
61+
) as Prefixed<ActionName>;
62+
}

‎js/plugins/googleai/src/embedder.ts

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ import {
2626
} from 'genkit';
2727
import { embedderRef } from 'genkit/embedder';
2828
import { embedder } from 'genkit/plugin';
29-
import { getApiKeyFromEnvVar } from './common.js';
30-
import type { PluginOptions } from './index.js';
29+
import { getApiKeyFromEnvVar, removePrefix } from './common';
30+
import type { PluginOptions } from './index';
3131

3232
export const TaskTypeSchema = z.enum([
3333
'RETRIEVAL_DOCUMENT',
@@ -117,9 +117,7 @@ export function defineGoogleAIEmbedder(
117117
);
118118
}
119119
// Extract the bare model name for lookup and API calls
120-
const apiModelName = name.startsWith('googleai/')
121-
? name.substring('googleai/'.length)
122-
: name;
120+
const apiModelName = removePrefix(name);
123121

124122
const embedderReference: EmbedderReference =
125123
SUPPORTED_MODELS[apiModelName] ??

‎js/plugins/googleai/src/gemini.ts

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,12 @@ import {
5959
import { downloadRequestMedia } from 'genkit/model/middleware';
6060
import { model } from 'genkit/plugin';
6161
import { runInNewSpan } from 'genkit/tracing';
62-
import { getApiKeyFromEnvVar, getGenkitClientHeader } from './common';
62+
import {
63+
ensurePrefixed,
64+
getApiKeyFromEnvVar,
65+
getGenkitClientHeader,
66+
removePrefix,
67+
} from './common';
6368
import { handleCacheIfNeeded } from './context-caching';
6469
import { extractCacheConfig } from './context-caching/utils';
6570

@@ -716,7 +721,7 @@ export function gemini(
716721
): ModelReference<typeof GeminiConfigSchema> {
717722
const nearestModel = nearestGeminiModelRef(version);
718723
return modelRef({
719-
name: `googleai/${version}`,
724+
name: ensurePrefixed(version),
720725
config: options,
721726
configSchema: GeminiConfigSchema,
722727
info: {
@@ -1150,9 +1155,7 @@ export function defineGoogleAIModel({
11501155
}
11511156

11521157
// Extract the API model name for lookup and API calls
1153-
const apiModelName = name.startsWith('googleai/')
1154-
? name.substring('googleai/'.length)
1155-
: name;
1158+
const apiModelName = removePrefix(name);
11561159

11571160
const modelReference: ModelReference<z.ZodTypeAny> =
11581161
SUPPORTED_GEMINI_MODELS[apiModelName] ??

‎js/plugins/googleai/src/imagen.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import {
2424
type ModelReference,
2525
} from 'genkit/model';
2626
import { model } from 'genkit/plugin';
27-
import { getApiKeyFromEnvVar } from './common.js';
27+
import { ensurePrefixed, getApiKeyFromEnvVar } from './common.js';
2828
import { predictModel } from './predict.js';
2929

3030
export type KNOWN_IMAGEN_MODELS = 'imagen-3.0-generate-002';
@@ -124,7 +124,7 @@ export function defineImagenModel(
124124
});
125125
}
126126
}
127-
const modelName = `googleai/${name}`;
127+
const modelName = ensurePrefixed(name);
128128
const modelReference: ModelReference<z.ZodTypeAny> = modelRef({
129129
name: modelName,
130130
info: {

‎js/plugins/googleai/src/index.ts

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import { logger } from 'genkit/logging';
2727
import { modelRef } from 'genkit/model';
2828
import { genkitPluginV2, type GenkitPluginV2 } from 'genkit/plugin';
2929
import type { ActionType } from 'genkit/registry';
30-
import { getApiKeyFromEnvVar } from './common.js';
30+
import { ensurePrefixed, getApiKeyFromEnvVar, removePrefix } from './common';
3131
import {
3232
SUPPORTED_MODELS as EMBEDDER_MODELS,
3333
GeminiEmbeddingConfigSchema,
@@ -227,7 +227,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 {
227227
Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => {
228228
actions.push(
229229
defineGoogleAIModel({
230-
name: `googleai/${name}`,
230+
name: ensurePrefixed(name),
231231
apiKey: options?.apiKey,
232232
apiVersion: 'v1beta',
233233
baseUrl: options?.baseUrl,
@@ -241,7 +241,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 {
241241
Object.keys(SUPPORTED_GEMINI_MODELS).forEach((name) => {
242242
actions.push(
243243
defineGoogleAIModel({
244-
name: `googleai/${name}`,
244+
name: ensurePrefixed(name),
245245
apiKey: options?.apiKey,
246246
baseUrl: options?.baseUrl,
247247
debugTraces: options?.experimental_debugTraces,
@@ -250,7 +250,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 {
250250
});
251251
Object.keys(EMBEDDER_MODELS).forEach((name) => {
252252
actions.push(
253-
defineGoogleAIEmbedder(`googleai/${name}`, {
253+
defineGoogleAIEmbedder(ensurePrefixed(name), {
254254
apiKey: options?.apiKey,
255255
})
256256
);
@@ -267,7 +267,7 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 {
267267
typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef;
268268
actions.push(
269269
defineGoogleAIModel({
270-
name: modelName,
270+
name: ensurePrefixed(modelName),
271271
apiKey: options?.apiKey,
272272
baseUrl: options?.baseUrl,
273273
info: { ...ref.info, label: `Google AI - ${modelName}` },
@@ -279,21 +279,24 @@ export function googleAIPlugin(options?: PluginOptions): GenkitPluginV2 {
279279
return actions;
280280
},
281281
async resolve(actionType: ActionType, actionName: string) {
282+
const rawActionName = removePrefix(actionName);
283+
const fullActionName = ensurePrefixed(rawActionName);
284+
282285
if (actionType === 'embedder') {
283-
return defineGoogleAIEmbedder(`googleai/${actionName}`, {
286+
return defineGoogleAIEmbedder(fullActionName, {
284287
apiKey: options?.apiKey,
285288
});
286-
} else if (actionName.startsWith('veo')) {
289+
} else if (rawActionName.startsWith('veo')) {
287290
if (actionType === 'background-model') {
288-
return defineVeoModel(`googleai/${actionName}`, options?.apiKey);
291+
return defineVeoModel(fullActionName, options?.apiKey);
289292
}
290293
} else if (actionType === 'model') {
291-
if (actionName.startsWith('imagen')) {
292-
return defineImagenModel(`googleai/${actionName}`, options?.apiKey);
294+
if (rawActionName.startsWith('imagen')) {
295+
return defineImagenModel(fullActionName, options?.apiKey);
293296
}
294297
// Fallback dynamic Gemini model
295298
return defineGoogleAIModel({
296-
name: `googleai/${actionName}`,
299+
name: fullActionName,
297300
apiKey: options?.apiKey,
298301
baseUrl: options?.baseUrl,
299302
debugTraces: options?.experimental_debugTraces,
@@ -339,22 +342,23 @@ export const googleAI = googleAIPlugin as GoogleAIPlugin;
339342
name: string,
340343
config?: any
341344
): ModelReference<z.ZodTypeAny> => {
345+
const fullModelName = ensurePrefixed(name);
342346
if (name.startsWith('imagen')) {
343347
return modelRef({
344-
name: `googleai/${name}`,
348+
name: fullModelName,
345349
config,
346350
configSchema: ImagenConfigSchema,
347351
});
348352
}
349353
if (name.startsWith('veo')) {
350354
return modelRef({
351-
name: `googleai/${name}`,
355+
name: fullModelName,
352356
config,
353357
configSchema: VeoConfigSchema,
354358
});
355359
}
356360
return modelRef({
357-
name: `googleai/${name}`,
361+
name: fullModelName,
358362
config,
359363
configSchema: GeminiConfigSchema,
360364
});
@@ -363,8 +367,9 @@ googleAI.embedder = (
363367
name: string,
364368
config?: GeminiEmbeddingConfig
365369
): EmbedderReference<typeof GeminiEmbeddingConfigSchema> => {
370+
const fullEmbedderName = ensurePrefixed(name);
366371
return embedderRef({
367-
name: `googleai/${name}`,
372+
name: fullEmbedderName,
368373
config,
369374
configSchema: GeminiEmbeddingConfigSchema,
370375
});

‎js/plugins/googleai/src/veo.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import {
2323
type ModelReference,
2424
} from 'genkit/model';
2525
import { backgroundModel } from 'genkit/plugin';
26-
import { getApiKeyFromEnvVar } from './common.js';
26+
import { ensurePrefixed, getApiKeyFromEnvVar } from './common.js';
2727
import { Operation as ApiOperation, checkOp, predictModel } from './predict.js';
2828

2929
export type KNOWN_VEO_MODELS = 'veo-2.0-generate-001';
@@ -136,7 +136,7 @@ export function defineVeoModel(
136136
});
137137
}
138138
}
139-
const modelName = `googleai/${name}`;
139+
const modelName = ensurePrefixed(name);
140140
const model: ModelReference<z.ZodTypeAny> = modelRef({
141141
name: modelName,
142142
info: {

0 commit comments

Comments
 (0)