refactor(plugins/checks): migrate to v2 API

js/plugins/checks/src/evaluation.ts

@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
-import { z, type EvaluatorAction, type Genkit } from 'genkit';
+import { z, type EvaluatorAction } from 'genkit';
 import type { BaseEvalDataPoint } from 'genkit/evaluator';
+import { evaluator } from 'genkit/plugin';
 import { runInNewSpan } from 'genkit/tracing';
 import type { GoogleAuth } from 'google-auth-library';
 import {
@@ -25,7 +26,6 @@ import {
 } from './metrics';
 
 export function checksEvaluators(
-  ai: Genkit,
   auth: GoogleAuth,
   metrics: ChecksEvaluationMetric[],
   projectId: string
@@ -42,7 +42,7 @@ export function checksEvaluators(
     }
   );
 
-  return createPolicyEvaluator(projectId, auth, ai, policy_configs);
+  return createPolicyEvaluator(projectId, auth, policy_configs);
 }
 
 const ResponseSchema = z.object({
@@ -58,10 +58,9 @@ const ResponseSchema = z.object({
 function createPolicyEvaluator(
   projectId: string,
   auth: GoogleAuth,
-  ai: Genkit,
   policy_config: ChecksEvaluationMetricConfig[]
 ): EvaluatorAction {
-  return ai.defineEvaluator(
+  return evaluator(
     {
       name: 'checks/guardrails',
       displayName: 'checks/guardrails',
@@ -83,7 +82,6 @@ function createPolicyEvaluator(
       };
 
       const response = await checksEvalInstance(
-        ai,
         projectId,
         auth,
         partialRequest,
@@ -109,14 +107,12 @@ function createPolicyEvaluator(
 }
 
 async function checksEvalInstance<ResponseType extends z.ZodTypeAny>(
-  ai: Genkit,
   projectId: string,
   auth: GoogleAuth,
   partialRequest: any,
   responseSchema: ResponseType
 ): Promise<z.infer<ResponseType>> {
   return await runInNewSpan(
-    ai,
     {
       metadata: {
         name: 'EvaluationService#evaluateInstances',

js/plugins/checks/src/index.ts

@@ -14,10 +14,10 @@
  * limitations under the License.
  */
 
-import type { Genkit } from 'genkit';
 import { logger } from 'genkit/logging';
 import type { ModelMiddleware } from 'genkit/model';
-import { genkitPlugin, type GenkitPlugin } from 'genkit/plugin';
+import { genkitPluginV2, type GenkitPluginV2 } from 'genkit/plugin';
+import { ActionType } from 'genkit/registry';
 import { GoogleAuth, type GoogleAuthOptions } from 'google-auth-library';
 import { checksEvaluators } from './evaluation.js';
 import {
@@ -47,23 +47,33 @@ const CHECKS_OAUTH_SCOPE = 'https://www.googleapis.com/auth/checks';
 /**
  * Add Google Checks evaluators.
  */
-export function checks(options?: PluginOptions): GenkitPlugin {
-  return genkitPlugin('checks', async (ai: Genkit) => {
-    const googleAuth = inititializeAuth(options?.googleAuthOptions);
+export function checks(options?: PluginOptions): GenkitPluginV2 {
+  const googleAuth = inititializeAuth(options?.googleAuthOptions);
 
-    const projectId = options?.projectId || (await googleAuth.getProjectId());
+  const projectId = options?.projectId || googleAuth.getProjectId();
 
-    if (!projectId) {
-      throw new Error(
-        `Checks Plugin is missing the 'projectId' configuration. Please set the 'GCLOUD_PROJECT' environment variable or explicitly pass 'projectId' into genkit config.`
-      );
-    }
+  if (!projectId) {
+    throw new Error(
+      `Checks Plugin is missing the 'projectId' configuration. Please set the 'GCLOUD_PROJECT' environment variable or explicitly pass 'projectId' into Genkit config.`
+    );
+  }
 
-    const metrics =
-      options?.evaluation && options.evaluation.metrics.length > 0
-        ? options.evaluation.metrics
-        : [];
-    checksEvaluators(ai, googleAuth, metrics, projectId);
+  const metrics =
+    options?.evaluation && options.evaluation.metrics.length > 0
+      ? options.evaluation.metrics
+      : [];
+
+  return genkitPluginV2({
+    name: 'checks',
+    init: async () => {
+      return [checksEvaluators(googleAuth, metrics, await projectId)];
+    },
+    resolve: async (actionType: ActionType, name: string) => {
+      return checksEvaluators(googleAuth, metrics, await projectId);
+    },
+    list: async () => {
+      return [checksEvaluators(googleAuth, metrics, await projectId).__action];
+    },
   });
 }