refactor(js/plugin/compat-oai): migrated compat-oai plugin to v2 plugin api

js/ai/src/embedder.ts

@@ -59,7 +59,7 @@ const EmbedRequestSchema = z.object({
 });
 
 export interface EmbedRequest<O = any> {
-  input: DocumentData[];
+  input: Document[];
   options?: O;
 }
 
@@ -120,7 +120,13 @@ export function embedder<ConfigSchema extends z.ZodTypeAny = z.ZodTypeAny>(
   ) => Promise<EmbedResponse>
 ) {
   const embedder = action(embedderActionOptions(options), (i, opts) =>
-    runner(i, opts)
+    runner(
+      {
+        input: i.input.map((dd) => new Document(dd)),
+        options: i.options,
+      },
+      opts
+    )
   );
   const ewm = withMetadata(
     embedder as Action<typeof EmbedRequestSchema, typeof EmbedResponseSchema>,

js/genkit/src/genkit.ts

@@ -125,6 +125,7 @@ import type { HasRegistry } from '@genkit-ai/core/registry';
 import type { BaseEvalDataPointSchema } from './evaluator.js';
 import { logger } from './logging.js';
 import {
+  ResolvableAction,
   isPluginV2,
   type GenkitPlugin,
   type GenkitPluginV2,
@@ -929,32 +930,17 @@ export class Genkit implements HasRegistry {
           name: plugin.name,
           async initializer() {
             logger.debug(`Initializing plugin ${plugin.name}:`);
+            if (!plugin.init) return;
+            const resolvedActions = await plugin.init();
+            resolvedActions?.forEach((resolvedAction) => {
+              registerActionV2(activeRegistry, resolvedAction, plugin);
+            });
           },
           async resolver(action: ActionType, target: string) {
+            if (!plugin.resolve) return;
             const resolvedAction = await plugin.resolve(action, target);
             if (resolvedAction) {
-              if (isBackgroundAction(resolvedAction)) {
-                registerBackgroundAction(activeRegistry, resolvedAction);
-              } else if (isAction(resolvedAction)) {
-                if (!resolvedAction.__action.actionType) {
-                  throw new GenkitError({
-                    status: 'INVALID_ARGUMENT',
-                    message:
-                      'Action type is missing for ' +
-                      resolvedAction.__action.name,
-                  });
-                }
-                activeRegistry.registerAction(
-                  resolvedAction.__action.actionType,
-                  resolvedAction
-                );
-              } else {
-                throw new GenkitError({
-                  status: 'INVALID_ARGUMENT',
-                  message:
-                    'Unkown action type returned from plugin ' + plugin.name,
-                });
-              }
+              registerActionV2(activeRegistry, resolvedAction, plugin);
             }
           },
           async listActions() {
@@ -998,6 +984,29 @@ export class Genkit implements HasRegistry {
   }
 }
 
+function registerActionV2(
+  registry: Registry,
+  resolvedAction: ResolvableAction,
+  plugin: GenkitPluginV2
+) {
+  if (isBackgroundAction(resolvedAction)) {
+    registerBackgroundAction(registry, resolvedAction);
+  } else if (isAction(resolvedAction)) {
+    if (!resolvedAction.__action.actionType) {
+      throw new GenkitError({
+        status: 'INVALID_ARGUMENT',
+        message: 'Action type is missing for ' + resolvedAction.__action.name,
+      });
+    }
+    registry.registerAction(resolvedAction.__action.actionType, resolvedAction);
+  } else {
+    throw new GenkitError({
+      status: 'INVALID_ARGUMENT',
+      message: 'Unkown action type returned from plugin ' + plugin.name,
+    });
+  }
+}
+
 /**
  * Initializes Genkit with a set of options.
  *

js/genkit/src/plugin.ts

@@ -37,7 +37,8 @@ export type ResolvableAction = Action | BackgroundAction;
 export interface GenkitPluginV2 {
   version: 'v2';
   name: string;
-  resolve: (
+  init?: () => ResolvableAction[] | Promise<ResolvableAction[]>;
+  resolve?: (
     actionType: ActionType,
     name: string
   ) => ResolvableAction | undefined | Promise<ResolvableAction | undefined>;

js/genkit/tests/plugins_test.ts

@@ -29,26 +29,16 @@ import {
 const v1Plugin = genkitPlugin(
   'myV1Plugin',
   (ai) => {
-    ai.defineModel(
-      {
-        name: 'myV1Plugin/model_eager',
-      },
-      async () => {
-        return {};
-      }
-    );
+    ai.defineModel({ name: 'myV1Plugin/model_eager' }, async () => {
+      return {};
+    });
   },
   async (ai, actionType, name) => {
     switch (actionType) {
       case 'model':
-        ai.defineModel(
-          {
-            name: 'myV1Plugin/' + name,
-          },
-          async () => {
-            return {};
-          }
-        );
+        ai.defineModel({ name: 'myV1Plugin/' + name }, async () => {
+          return {};
+        });
       case 'background-model':
         ai.defineBackgroundModel({
           name: 'myV1Plugin/' + name,
@@ -78,17 +68,19 @@ const v1Plugin = genkitPlugin(
 
 const v2Plugin = genkitPluginV2({
   name: 'myV2Plugin',
+  init() {
+    return [
+      model({ name: 'myV2Plugin/model_eager' }, async () => {
+        return {};
+      }),
+    ];
+  },
   resolve(actionType, name) {
     switch (actionType) {
       case 'model':
-        return model(
-          {
-            name: 'myV2Plugin/' + name,
-          },
-          async () => {
-            return {};
-          }
-        );
+        return model({ name: 'myV2Plugin/' + name }, async () => {
+          return {};
+        });
       case 'background-model':
         return backgroundModel({
           name: 'myV2Plugin/' + name,
@@ -139,6 +131,7 @@ describe('session', () => {
         '/embedder/myV2Plugin/potential_embedder',
         '/model/myV1Plugin/potential_model',
         '/model/myV1Plugin/model_eager',
+        '/model/myV2Plugin/model_eager',
         '/model/myV2Plugin/potential_model',
       ])
     );
@@ -148,7 +141,10 @@ describe('session', () => {
           (a) => a.startsWith('/model') || a.startsWith('/embedder')
         )
       ),
-      new Set(['/model/myV1Plugin/model_eager'])
+      new Set([
+        '/model/myV1Plugin/model_eager',
+        '/model/myV2Plugin/model_eager',
+      ])
     );
   });
 
@@ -168,6 +164,7 @@ describe('session', () => {
       ),
       new Set([
         '/model/myV1Plugin/model_eager',
+        '/model/myV2Plugin/model_eager',
         '/model/myV1Plugin/potential_model',
       ])
     );
@@ -214,6 +211,7 @@ describe('session', () => {
       ),
       new Set([
         '/model/myV1Plugin/model_eager',
+        '/model/myV2Plugin/model_eager',
         '/model/myV2Plugin/potential_model',
       ])
     );

js/plugins/compat-oai/src/audio.ts

@@ -17,11 +17,11 @@
 import type {
   GenerateRequest,
   GenerateResponseData,
-  Genkit,
   ModelReference,
 } from 'genkit';
 import { GenerationCommonConfigSchema, Message, modelRef, z } from 'genkit';
 import type { ModelAction, ModelInfo } from 'genkit/model';
+import { model } from 'genkit/plugin';
 import type OpenAI from 'openai';
 import { Response } from 'openai/core.mjs';
 import type {
@@ -181,19 +181,17 @@ async function toGenerateResponse(
 export function defineCompatOpenAISpeechModel<
   CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
 >(params: {
-  ai: Genkit;
   name: string;
   client: OpenAI;
   modelRef?: ModelReference<CustomOptions>;
   requestBuilder?: SpeechRequestBuilder;
 }): ModelAction {
-  const { ai, name, client, modelRef, requestBuilder } = params;
+  const { name, client, modelRef, requestBuilder } = params;
   const modelName = name.substring(name.indexOf('/') + 1);
 
-  return ai.defineModel(
+  return model(
     {
       name,
-      apiVersion: 'v2',
       ...modelRef?.info,
       configSchema: modelRef?.configSchema,
     },
@@ -335,18 +333,16 @@ function transcriptionToGenerateResponse(
 export function defineCompatOpenAITranscriptionModel<
   CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
 >(params: {
-  ai: Genkit;
   name: string;
   client: OpenAI;
   modelRef?: ModelReference<CustomOptions>;
   requestBuilder?: TranscriptionRequestBuilder;
 }): ModelAction {
-  const { ai, name, client, modelRef, requestBuilder } = params;
+  const { name, client, modelRef, requestBuilder } = params;
 
-  return ai.defineModel(
+  return model(
     {
       name,
-      apiVersion: 'v2',
       ...modelRef?.info,
       configSchema: modelRef?.configSchema,
     },

js/plugins/compat-oai/src/deepseek/index.ts

@@ -16,14 +16,13 @@
 
 import {
   ActionMetadata,
-  Genkit,
   GenkitError,
   modelActionMetadata,
   ModelReference,
   z,
 } from 'genkit';
 import { logger } from 'genkit/logging';
-import { GenkitPlugin } from 'genkit/plugin';
+import { GenkitPluginV2 } from 'genkit/plugin';
 import { ActionType } from 'genkit/registry';
 import OpenAI from 'openai';
 import { openAICompatible, PluginOptions } from '../index.js';
@@ -38,7 +37,6 @@ import {
 export type DeepSeekPluginOptions = Omit<PluginOptions, 'name' | 'baseURL'>;
 
 const resolver = async (
-  ai: Genkit,
   client: OpenAI,
   actionType: ActionType,
   actionName: string
@@ -47,8 +45,7 @@ const resolver = async (
     const modelRef = deepSeekModelRef({
       name: `deepseek/${actionName}`,
     });
-    defineCompatOpenAIModel({
-      ai,
+    return defineCompatOpenAIModel({
       name: modelRef.name,
       client,
       modelRef,
@@ -57,6 +54,7 @@ const resolver = async (
   } else {
     logger.warn('Only model actions are supported by the DeepSeek plugin');
   }
+  return undefined;
 };
 
 const listActions = async (client: OpenAI): Promise<ActionMetadata[]> => {
@@ -78,7 +76,9 @@ const listActions = async (client: OpenAI): Promise<ActionMetadata[]> => {
   );
 };
 
-export function deepSeekPlugin(options?: DeepSeekPluginOptions): GenkitPlugin {
+export function deepSeekPlugin(
+  options?: DeepSeekPluginOptions
+): GenkitPluginV2 {
   const apiKey = options?.apiKey ?? process.env.DEEPSEEK_API_KEY;
   if (!apiKey) {
     throw new GenkitError({
@@ -92,10 +92,9 @@ export function deepSeekPlugin(options?: DeepSeekPluginOptions): GenkitPlugin {
     baseURL: 'https://api.deepseek.com',
     apiKey,
     ...options,
-    initializer: async (ai, client) => {
-      Object.values(SUPPORTED_DEEPSEEK_MODELS).forEach((modelRef) =>
+    initializer: async (client) => {
+      return Object.values(SUPPORTED_DEEPSEEK_MODELS).map((modelRef) =>
         defineCompatOpenAIModel({
-          ai,
           name: modelRef.name,
           client,
           modelRef,
@@ -109,7 +108,7 @@ export function deepSeekPlugin(options?: DeepSeekPluginOptions): GenkitPlugin {
 }
 
 export type DeepSeekPlugin = {
-  (params?: DeepSeekPluginOptions): GenkitPlugin;
+  (params?: DeepSeekPluginOptions): GenkitPluginV2;
   model(
     name: keyof typeof SUPPORTED_DEEPSEEK_MODELS,
     config?: z.infer<typeof DeepSeekChatCompletionConfigSchema>

js/plugins/compat-oai/src/embedder.ts

@@ -17,7 +17,8 @@
 
 // import { defineEmbedder, embedderRef } from '@genkit-ai/ai/embedder';
 
-import type { EmbedderAction, EmbedderReference, Genkit } from 'genkit';
+import type { EmbedderAction, EmbedderReference } from 'genkit';
+import { embedder } from 'genkit/plugin';
 import OpenAI from 'openai';
 
 /**
@@ -34,25 +35,24 @@ import OpenAI from 'openai';
  * @returns the created {@link EmbedderAction}
  */
 export function defineCompatOpenAIEmbedder(params: {
-  ai: Genkit;
   name: string;
   client: OpenAI;
   embedderRef?: EmbedderReference;
 }): EmbedderAction {
-  const { ai, name, client, embedderRef } = params;
+  const { name, client, embedderRef } = params;
   const modelName = name.substring(name.indexOf('/') + 1);
 
-  return ai.defineEmbedder(
+  return embedder(
     {
       name,
       configSchema: embedderRef?.configSchema,
       ...embedderRef?.info,
     },
-    async (input, options) => {
-      const { encodingFormat: encoding_format, ...restOfConfig } = options;
+    async (req) => {
+      const { encodingFormat: encoding_format, ...restOfConfig } = req.options;
       const embeddings = await client.embeddings.create({
         model: modelName!,
-        input: input.map((d) => d.text),
+        input: req.input.map((d) => d.text),
         encoding_format,
         ...restOfConfig,
       });