refactor(js/plugins/vertexai): extract common logic

js/plugins/vertexai/package.json

@@ -67,5 +67,8 @@
       "types": "./lib/index.d.ts",
       "default": "./lib/index.js"
     }
+  },
+  "typesVersions": {
+    "*": {}
   }
 }

js/plugins/vertexai/src/common/auth.ts

@@ -0,0 +1,48 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { GoogleAuth, GoogleAuthOptions } from 'google-auth-library';
+
+const CLOUD_PLATFROM_OAUTH_SCOPE =
+  'https://www.googleapis.com/auth/cloud-platform';
+
+/**
+ * Authenticate with Google Cloud
+ * @param authOptions The authentication options to use.
+ * @returns The GoogleAuth object.
+ */
+export const authenticate: (authOptions?: GoogleAuthOptions) => GoogleAuth = (
+  authOptions?: GoogleAuthOptions
+) => {
+  let authClient;
+
+  // Allow customers to pass in cloud credentials from environment variables
+  // following: https://github.com/googleapis/google-auth-library-nodejs?tab=readme-ov-file#loading-credentials-from-environment-variables
+  if (process.env.GCLOUD_SERVICE_ACCOUNT_CREDS) {
+    const serviceAccountCreds = JSON.parse(
+      process.env.GCLOUD_SERVICE_ACCOUNT_CREDS
+    );
+    authOptions = {
+      credentials: serviceAccountCreds,
+      scopes: [CLOUD_PLATFROM_OAUTH_SCOPE],
+    };
+    authClient = new GoogleAuth(authOptions);
+  } else {
+    authClient = new GoogleAuth(
+      authOptions ?? { scopes: [CLOUD_PLATFROM_OAUTH_SCOPE] }
+    );
+  }
+
+  return authClient;
+};

js/plugins/vertexai/src/common/global.ts

@@ -0,0 +1,21 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+export const DEFAULT_LOCATION = 'us-central1';
+
+export const confError = (parameter: string, envVariableName: string) => {
+  return new Error(
+    `VertexAI Plugin is missing the '${parameter}' configuration. Please set the '${envVariableName}' environment variable or explicitly pass '${parameter}' into genkit config.`
+  );
+};

js/plugins/vertexai/src/common/types.ts

@@ -0,0 +1,24 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { GoogleAuthOptions } from 'google-auth-library';
+
+export interface BasePluginOptions {
+  /** The Google Cloud project id to call. */
+  projectId?: string;
+  /** The Google Cloud region to call. */
+  location: string;
+  /** Provide custom authentication configuration for connecting to Vertex AI. */
+  googleAuth?: GoogleAuthOptions;
+}

js/plugins/vertexai/src/index.ts

@@ -18,7 +18,7 @@ import { VertexAI } from '@google-cloud/vertexai';
 import { genkitPlugin, Plugin, z } from 'genkit';
 import { GenerateRequest, ModelReference } from 'genkit/model';
 import { IndexerAction, RetrieverAction } from 'genkit/retriever';
-import { GoogleAuth, GoogleAuthOptions } from 'google-auth-library';
+import { GoogleAuthOptions } from 'google-auth-library';
 import {
   anthropicModel,
   claude35Sonnet,
@@ -27,6 +27,8 @@ import {
   claude3Sonnet,
   SUPPORTED_ANTHROPIC_MODELS,
 } from './anthropic.js';
+import { authenticate } from './common/auth.js';
+import { confError, DEFAULT_LOCATION } from './common/global.js';
 import {
   SUPPORTED_EMBEDDER_MODELS,
   textEmbedding004,
@@ -140,43 +142,19 @@ export interface PluginOptions {
   rerankOptions?: VertexRerankerConfig[];
 }
 
-const CLOUD_PLATFROM_OAUTH_SCOPE =
-  'https://www.googleapis.com/auth/cloud-platform';
-
 /**
  * Add Google Cloud Vertex AI to Genkit. Includes Gemini and Imagen models and text embedder.
  */
 export const vertexAI: Plugin<[PluginOptions] | []> = genkitPlugin(
   'vertexai',
   async (options?: PluginOptions) => {
-    let authClient;
-    let authOptions = options?.googleAuth;
-
-    // Allow customers to pass in cloud credentials from environment variables
-    // following: https://github.com/googleapis/google-auth-library-nodejs?tab=readme-ov-file#loading-credentials-from-environment-variables
-    if (process.env.GCLOUD_SERVICE_ACCOUNT_CREDS) {
-      const serviceAccountCreds = JSON.parse(
-        process.env.GCLOUD_SERVICE_ACCOUNT_CREDS
-      );
-      authOptions = {
-        credentials: serviceAccountCreds,
-        scopes: [CLOUD_PLATFROM_OAUTH_SCOPE],
-      };
-      authClient = new GoogleAuth(authOptions);
-    } else {
-      authClient = new GoogleAuth(
-        authOptions ?? { scopes: [CLOUD_PLATFROM_OAUTH_SCOPE] }
-      );
-    }
+    // Authenticate with Google Cloud
+    const authOptions = options?.googleAuth;
+    const authClient = authenticate(authOptions);
 
     const projectId = options?.projectId || (await authClient.getProjectId());
+    const location = options?.location || DEFAULT_LOCATION;
 
-    const location = options?.location || 'us-central1';
-    const confError = (parameter: string, envVariableName: string) => {
-      return new Error(
-        `VertexAI Plugin is missing the '${parameter}' configuration. Please set the '${envVariableName}' environment variable or explicitly pass '${parameter}' into genkit config.`
-      );
-    };
     if (!location) {
       throw confError('location', 'GCLOUD_LOCATION');
     }