Skip to content

If dynamic2 #2827

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

If dynamic2 #2827

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d925df6
feat(js/plugins/vertexai): Dynamic model support for Gemini models
Apr 23, 2025
d2286fd
test(js/plugins/vertexai): modified the test
Apr 25, 2025
8861f98
feat(js/plugins/googleai) Dynamic model support for Gemini models
Apr 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions js/core/src/registry.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,9 +241,9 @@ export class Registry {
}
return cached;
},
resolver: async (actionType: ActionType, target: string) => {
resolver: async (actionType: ActionType, actionName: string) => {
if (provider.resolver) {
await provider.resolver(actionType, target);
await provider.resolver(actionType, actionName);
}
},
};
Expand All @@ -261,19 +261,20 @@ export class Registry {
/**
* Resolves a new Action dynamically by registering it.
* @param pluginName The name of the plugin
* @param action
* @param actionType The type of the action
* @param actionName The name of the action
* @returns
*/
async resolvePluginAction(
pluginName: string,
action: ActionType,
target: string
actionType: ActionType,
actionName: string
) {
const plugin = this.pluginsByName[pluginName];
if (plugin) {
return await runOutsideActionRuntimeContext(this, async () => {
if (plugin.resolver) {
await plugin.resolver(action, target);
await plugin.resolver(actionType, actionName);
}
});
}
Expand Down
1 change: 1 addition & 0 deletions js/plugins/googleai/src/gemini.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,7 @@ function nearestGeminiModelRef(
version,
});
}

return GENERIC_GEMINI_MODEL.withConfig({ ...options, version });
}

Expand Down
188 changes: 117 additions & 71 deletions js/plugins/googleai/src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import { Genkit, ModelReference } from 'genkit';
import { GenkitPlugin, genkitPlugin } from 'genkit/plugin';
import { ActionType } from 'genkit/registry';
import {
SUPPORTED_MODELS as EMBEDDER_MODELS,
defineGoogleAIEmbedder,
Expand Down Expand Up @@ -81,82 +82,127 @@ export interface PluginOptions {
experimental_debugTraces?: boolean;
}

/**
* Google Gemini Developer API plugin.
*/
export function googleAI(options?: PluginOptions): GenkitPlugin {
return genkitPlugin('googleai', async (ai: Genkit) => {
let apiVersions = ['v1'];
async function initializer(ai: Genkit, options?: PluginOptions) {
let apiVersions = ['v1'];

if (options?.apiVersion) {
if (Array.isArray(options?.apiVersion)) {
apiVersions = options?.apiVersion;
} else {
apiVersions = [options?.apiVersion];
}
if (options?.apiVersion) {
if (Array.isArray(options?.apiVersion)) {
apiVersions = options?.apiVersion;
} else {
apiVersions = [options?.apiVersion];
}
}

if (apiVersions.includes('v1beta')) {
Object.keys(SUPPORTED_V15_MODELS).forEach((name) =>
defineGoogleAIModel({
ai,
name,
apiKey: options?.apiKey,
apiVersion: 'v1beta',
baseUrl: options?.baseUrl,
debugTraces: options?.experimental_debugTraces,
})
);
}
if (apiVersions.includes('v1')) {
Object.keys(SUPPORTED_V1_MODELS).forEach((name) =>
defineGoogleAIModel({
ai,
name,
apiKey: options?.apiKey,
apiVersion: undefined,
baseUrl: options?.baseUrl,
debugTraces: options?.experimental_debugTraces,
})
);
Object.keys(SUPPORTED_V15_MODELS).forEach((name) =>
defineGoogleAIModel({
ai,
name,
apiKey: options?.apiKey,
apiVersion: undefined,
baseUrl: options?.baseUrl,
debugTraces: options?.experimental_debugTraces,
})
);
Object.keys(EMBEDDER_MODELS).forEach((name) =>
defineGoogleAIEmbedder(ai, name, { apiKey: options?.apiKey })
);
}
if (apiVersions.includes('v1beta')) {
Object.keys(SUPPORTED_V15_MODELS).forEach((name) =>
defineGoogleAIModel({
ai,
name,
apiKey: options?.apiKey,
apiVersion: 'v1beta',
baseUrl: options?.baseUrl,
debugTraces: options?.experimental_debugTraces,
})
);
}
if (apiVersions.includes('v1')) {
Object.keys(SUPPORTED_V1_MODELS).forEach((name) =>
defineGoogleAIModel({
ai,
name,
apiKey: options?.apiKey,
apiVersion: undefined,
baseUrl: options?.baseUrl,
debugTraces: options?.experimental_debugTraces,
})
);
Object.keys(SUPPORTED_V15_MODELS).forEach((name) =>
defineGoogleAIModel({
ai,
name,
apiKey: options?.apiKey,
apiVersion: undefined,
baseUrl: options?.baseUrl,
debugTraces: options?.experimental_debugTraces,
})
);
Object.keys(EMBEDDER_MODELS).forEach((name) =>
defineGoogleAIEmbedder(ai, name, { apiKey: options?.apiKey })
);
}

if (options?.models) {
for (const modelOrRef of options?.models) {
const modelName =
typeof modelOrRef === 'string'
? modelOrRef
: // strip out the `googleai/` prefix
modelOrRef.name.split('/')[1];
const modelRef =
typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef;
defineGoogleAIModel({
ai,
name: modelName,
apiKey: options?.apiKey,
baseUrl: options?.baseUrl,
info: {
...modelRef.info,
label: `Google AI - ${modelName}`,
},
debugTraces: options?.experimental_debugTraces,
});
}
if (options?.models) {
for (const modelOrRef of options?.models) {
const modelName =
typeof modelOrRef === 'string'
? modelOrRef
: // strip out the `googleai/` prefix
modelOrRef.name.split('/')[1];
const modelRef =
typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef;
defineGoogleAIModel({
ai,
name: modelName,
apiKey: options?.apiKey,
baseUrl: options?.baseUrl,
info: {
...modelRef.info,
label: `Google AI - ${modelName}`,
},
debugTraces: options?.experimental_debugTraces,
});
}
});
}
}

async function resolver(
ai: Genkit,
actionType: ActionType,
actionName: string,
options?: PluginOptions
) {
// TODO: also support other actions like 'embedder'
switch (actionType) {
case 'model':
await resolveModel(ai, actionName, options);
break;
default:
// no-op
}
}

async function resolveModel(
ai: Genkit,
actionName: string,
options?: PluginOptions
) {
if (actionName.includes('gemini')) {
const modelRef = gemini(actionName);
defineGoogleAIModel({
ai,
name: modelRef.name,
apiKey: options?.apiKey,
baseUrl: options?.baseUrl,
info: {
...modelRef.info,
label: `Google AI - ${actionName}`,
},
debugTraces: options?.experimental_debugTraces,
});
}
// TODO: Support other models
}

/**
* Google Gemini Developer API plugin.
*/
export function googleAI(options?: PluginOptions): GenkitPlugin {
return genkitPlugin(
'googleai',
async (ai: Genkit) => await initializer(ai, options),
async (ai: Genkit, actionType: ActionType, actionName: string) =>
await resolver(ai, actionType, actionName, options)
);
}

export default googleAI;
20 changes: 13 additions & 7 deletions js/plugins/googleai/tests/gemini_test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -463,7 +463,7 @@ describe('plugin', () => {
assert.strictEqual(flash.__action.name, 'googleai/gemini-1.5-flash');
});

it('references explicitly registered models', async () => {
it('references both pre-registered or dynamic models', async () => {
const flash002Ref = gemini('gemini-1.5-flash-002');
const ai = genkit({
plugins: [
Expand Down Expand Up @@ -525,13 +525,19 @@ describe('plugin', () => {
GENERIC_GEMINI_MODEL.info! // <---- generic model fallback
);

// this one is not registered
const flash003Ref = gemini('gemini-1.5-flash-003');
assert.strictEqual(flash003Ref.name, 'googleai/gemini-1.5-flash-003');
const flash003 = await ai.registry.lookupAction(
`/model/${flash003Ref.name}`
// this one is dynamically resolved (not pre-registered)
const giraffeRef = gemini('gemini-4.5-giraffe');
assert.strictEqual(giraffeRef.name, 'googleai/gemini-4.5-giraffe');
const giraffe = await ai.registry.lookupAction(
`/model/${giraffeRef.name}`
);
assert.ok(giraffe);
assert.strictEqual(giraffe.__action.name, 'googleai/gemini-4.5-giraffe');
assertEqualModelInfo(
giraffe.__action.metadata?.model,
'Google AI - gemini-4.5-giraffe',
GENERIC_GEMINI_MODEL.info! // <---- generic model fallback
);
assert.ok(flash003 === undefined);
});
});
});
Expand Down
Loading