feat(js/plugins/compat-oai): Support dynamic model resolution (#3117)

ssbushi · web-flow · commit b614b48086c8 · 2025-06-23T09:48:28.000-04:00
diff --git a/js/plugins/compat-oai/src/index.ts b/js/plugins/compat-oai/src/index.ts
@@ -14,13 +14,21 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { type Genkit } from 'genkit';
+import { ActionMetadata, type Genkit } from 'genkit';
 import { genkitPlugin } from 'genkit/plugin';
+import { ActionType } from 'genkit/registry';
 import { OpenAI, type ClientOptions } from 'openai';
 
 export interface PluginOptions extends Partial<ClientOptions> {
   name: string;
   initializer?: (ai: Genkit, client: OpenAI) => Promise<void>;
+  resolver?: (
+    ai: Genkit,
+    client: OpenAI,
+    actionType: ActionType,
+    actionName: string
+  ) => Promise<void>;
+  listActions?: (client: OpenAI) => Promise<ActionMetadata[]>;
 }
 
 /**
@@ -75,12 +83,29 @@ export interface PluginOptions extends Partial<ClientOptions> {
  * });
  * ```
  */
-export const openAICompatible = (options: PluginOptions) =>
-  genkitPlugin(options.name, async (ai: Genkit) => {
-    const client = new OpenAI(options);
-    if (options.initializer) {
-      await options.initializer(ai, client);
+export const openAICompatible = (options: PluginOptions) => {
+  const client = new OpenAI(options);
+  let listActionsCache;
+  return genkitPlugin(
+    options.name,
+    async (ai: Genkit) => {
+      if (options.initializer) {
+        await options.initializer(ai, client);
+      }
+    },
+    async (ai: Genkit, actionType: ActionType, actionName: string) => {
+      if (options.resolver) {
+        await options.resolver(ai, client, actionType, actionName);
+      }
+    },
+    async () => {
+      if (options.listActions) {
+        if (listActionsCache) return listActionsCache;
+        listActionsCache = await options.listActions(client);
+        return listActionsCache;
+      }
     }
-  });
+  );
+};
 
 export default openAICompatible;
diff --git a/js/plugins/compat-oai/src/openai/dalle.ts b/js/plugins/compat-oai/src/openai/dalle.ts
@@ -15,9 +15,9 @@
  * limitations under the License.
  */
 import { z } from 'genkit';
-import { modelRef } from 'genkit/model';
+import { ModelInfo, modelRef } from 'genkit/model';
 
-export const DallE3ConfigSchema = z.object({
+export const ImageGenerationConfigSchema = z.object({
   size: z.enum(['1024x1024', '1792x1024', '1024x1792']).optional(),
   style: z.enum(['vivid', 'natural']).optional(),
   user: z.string().optional(),
@@ -26,19 +26,23 @@ export const DallE3ConfigSchema = z.object({
   response_format: z.enum(['b64_json', 'url']).optional(),
 });
 
+export const IMAGE_GENERATION_MODEL_INFO: ModelInfo = {
+  supports: {
+    media: false,
+    output: ['media'],
+    multiturn: false,
+    systemRole: false,
+    tools: false,
+  },
+};
+
 export const dallE3 = modelRef({
   name: 'openai/dall-e-3',
   info: {
     label: 'OpenAI - DALL-E 3',
-    supports: {
-      media: false,
-      output: ['media'],
-      multiturn: false,
-      systemRole: false,
-      tools: false,
-    },
+    ...IMAGE_GENERATION_MODEL_INFO,
   },
-  configSchema: DallE3ConfigSchema,
+  configSchema: ImageGenerationConfigSchema,
 });
 
 export const SUPPORTED_IMAGE_MODELS = {
diff --git a/js/plugins/compat-oai/src/openai/index.ts b/js/plugins/compat-oai/src/openai/index.ts
@@ -15,6 +15,20 @@
  * limitations under the License.
  */
 
+import {
+  ActionMetadata,
+  embedderActionMetadata,
+  embedderRef,
+  EmbedderReference,
+  Genkit,
+  modelActionMetadata,
+  modelRef,
+  ModelReference,
+  z,
+} from 'genkit';
+import { GenkitPlugin } from 'genkit/plugin';
+import { ActionType } from 'genkit/registry';
+import OpenAI from 'openai';
 import {
   defineCompatOpenAISpeechModel,
   defineCompatOpenAITranscriptionModel,
@@ -23,14 +37,104 @@ import { defineCompatOpenAIEmbedder } from '../embedder.js';
 import { defineCompatOpenAIImageModel } from '../image.js';
 import openAICompatible, { PluginOptions } from '../index.js';
 import { defineCompatOpenAIModel } from '../model.js';
-import { SUPPORTED_IMAGE_MODELS } from './dalle.js';
-import { SUPPORTED_EMBEDDING_MODELS } from './embedder.js';
-import { SUPPORTED_GPT_MODELS } from './gpt.js';
-import { SUPPORTED_TTS_MODELS } from './tts.js';
-import { SUPPORTED_STT_MODELS } from './whisper.js';
+import {
+  IMAGE_GENERATION_MODEL_INFO,
+  ImageGenerationConfigSchema,
+  SUPPORTED_IMAGE_MODELS,
+} from './dalle.js';
+import {
+  SUPPORTED_EMBEDDING_MODELS,
+  TextEmbeddingConfigSchema,
+} from './embedder.js';
+import { ChatCompletionConfigSchema, SUPPORTED_GPT_MODELS } from './gpt.js';
+import {
+  SPEECH_MODEL_INFO,
+  SpeechConfigSchema,
+  SUPPORTED_TTS_MODELS,
+} from './tts.js';
+import { SUPPORTED_STT_MODELS, TranscriptionConfigSchema } from './whisper.js';
 
 export type OpenAIPluginOptions = Exclude<PluginOptions, 'name'>;
 
+const resolver = async (
+  ai: Genkit,
+  client: OpenAI,
+  actionType: ActionType,
+  actionName: string
+) => {
+  if (actionType === 'embedder') {
+    defineCompatOpenAIEmbedder({ ai, name: `openai/${actionName}`, client });
+  } else if (
+    actionName.includes('gpt-image-1') ||
+    actionName.includes('dall-e')
+  ) {
+    defineCompatOpenAIImageModel({ ai, name: `openai/${actionName}`, client });
+  } else if (actionName.includes('tts')) {
+    defineCompatOpenAISpeechModel({ ai, name: `openai/${actionName}`, client });
+  } else if (
+    actionName.includes('whisper') ||
+    actionName.includes('transcribe')
+  ) {
+    defineCompatOpenAITranscriptionModel({
+      ai,
+      name: `openai/${actionName}`,
+      client,
+    });
+  } else {
+    defineCompatOpenAIModel({
+      ai,
+      name: `openai/${actionName}`,
+      client,
+    });
+  }
+};
+
+const listActions = async (client: OpenAI): Promise<ActionMetadata[]> => {
+  return await client.models.list().then((response) =>
+    response.data
+      .filter((model) => model.object === 'model')
+      .map((model: OpenAI.Model) => {
+        if (model.id.includes('embedding')) {
+          return embedderActionMetadata({
+            name: `openai/${model.id}`,
+            configSchema: TextEmbeddingConfigSchema,
+            info: SUPPORTED_EMBEDDING_MODELS[model.id]?.info,
+          });
+        } else if (
+          model.id.includes('gpt-image-1') ||
+          model.id.includes('dall-e')
+        ) {
+          return modelActionMetadata({
+            name: `openai/${model.id}`,
+            configSchema: ImageGenerationConfigSchema,
+            info: IMAGE_GENERATION_MODEL_INFO,
+          });
+        } else if (model.id.includes('tts')) {
+          return modelActionMetadata({
+            name: `openai/${model.id}`,
+            configSchema: SpeechConfigSchema,
+            info: SPEECH_MODEL_INFO,
+          });
+        } else if (
+          model.id.includes('whisper') ||
+          model.id.includes('transcribe')
+        ) {
+          return modelActionMetadata({
+            name: `openai/${model.id}`,
+            configSchema: TranscriptionConfigSchema,
+            info: SPEECH_MODEL_INFO,
+          });
+        } else {
+          return modelActionMetadata({
+            name: `openai/${model.id}`,
+            configSchema: ChatCompletionConfigSchema,
+            info: SUPPORTED_GPT_MODELS[model.id]?.info,
+          });
+        }
+      })
+  );
+};
+
 /**
  * This module provides an interface to the OpenAI models through the Genkit
  * plugin system. It allows users to interact with various models by providing
@@ -60,8 +164,8 @@ export type OpenAIPluginOptions = Exclude<PluginOptions, 'name'>;
  * });
  * ```
  */
-export const openAI = (options?: OpenAIPluginOptions) =>
-  openAICompatible({
+export function openAIPlugin(options?: OpenAIPluginOptions): GenkitPlugin {
+  return openAICompatible({
     name: 'openai',
     ...options,
     initializer: async (ai, client) => {
@@ -101,6 +205,59 @@ export const openAI = (options?: OpenAIPluginOptions) =>
         })
       );
     },
+    resolver,
+    listActions,
+  });
+}
+
+export type OpenAIPlugin = {
+  (params?: OpenAIPluginOptions): GenkitPlugin;
+  model(name: string, config?: any): ModelReference<z.ZodTypeAny>;
+  embedder(name: string, config?: any): EmbedderReference<z.ZodTypeAny>;
+};
+
+export const openAI = openAIPlugin as OpenAIPlugin;
+// provide generic implementation for the model function overloads.
+(openAI as any).model = (
+  name: string,
+  config?: any
+): ModelReference<z.ZodTypeAny> => {
+  if (name.includes('gpt-image-1') || name.includes('dall-e')) {
+    return modelRef({
+      name: `openai/${name}`,
+      config,
+      configSchema: ImageGenerationConfigSchema,
+    });
+  }
+  if (name.includes('tts')) {
+    return modelRef({
+      name: `openai/${name}`,
+      config,
+      configSchema: SpeechConfigSchema,
+    });
+  }
+  if (name.includes('whisper') || name.includes('transcribe')) {
+    return modelRef({
+      name: `openai/${name}`,
+      config,
+      configSchema: TranscriptionConfigSchema,
+    });
+  }
+  return modelRef({
+    name: `openai/${name}`,
+    config,
+    configSchema: ChatCompletionConfigSchema,
+  });
+};
+openAI.embedder = (
+  name: string,
+  config?: any
+): EmbedderReference<z.ZodTypeAny> => {
+  return embedderRef({
+    name: `openai/${name}`,
+    config,
+    configSchema: TextEmbeddingConfigSchema,
   });
+};
 
 export default openAI;
diff --git a/js/plugins/compat-oai/src/openai/tts.ts b/js/plugins/compat-oai/src/openai/tts.ts
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 import { z } from 'genkit';
-import { modelRef } from 'genkit/model';
+import { ModelInfo, modelRef } from 'genkit/model';
 
 export const SpeechConfigSchema = z.object({
   voice: z
@@ -27,17 +27,21 @@ export const SpeechConfigSchema = z.object({
     .optional(),
 });
 
+export const SPEECH_MODEL_INFO: ModelInfo = {
+  supports: {
+    media: false,
+    output: ['media'],
+    multiturn: false,
+    systemRole: false,
+    tools: false,
+  },
+};
+
 export const tts1 = modelRef({
   name: 'openai/tts-1',
   info: {
     label: 'OpenAI - Text-to-speech 1',
-    supports: {
-      media: false,
-      output: ['media'],
-      multiturn: false,
-      systemRole: false,
-      tools: false,
-    },
+    ...SPEECH_MODEL_INFO,
   },
   configSchema: SpeechConfigSchema,
 });
@@ -46,13 +50,7 @@ export const tts1Hd = modelRef({
   name: 'openai/tts-1-hd',
   info: {
     label: 'OpenAI - Text-to-speech 1 HD',
-    supports: {
-      media: false,
-      output: ['media'],
-      multiturn: false,
-      systemRole: false,
-      tools: false,
-    },
+    ...SPEECH_MODEL_INFO,
   },
   configSchema: SpeechConfigSchema,
 });
@@ -61,13 +59,7 @@ export const gpt4oMiniTts = modelRef({
   name: 'openai/gpt-4o-mini-tts',
   info: {
     label: 'OpenAI - GPT-4o Mini Text-to-speech',
-    supports: {
-      media: false,
-      output: ['media'],
-      multiturn: false,
-      systemRole: false,
-      tools: false,
-    },
+    ...SPEECH_MODEL_INFO,
   },
   configSchema: SpeechConfigSchema.omit({ speed: true }),
 });
diff --git a/js/plugins/compat-oai/src/openai/whisper.ts b/js/plugins/compat-oai/src/openai/whisper.ts
diff --git a/js/testapps/compat-oai/src/index.ts b/js/testapps/compat-oai/src/index.ts
