fix(js/plugin): added init function to plugin v2 (#3431)

Author: pavelgj

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',
       ])
     );