feat: include all supported authentication options

PiperOrigin-RevId: 595200945
googleapis · Jan 3, 2024 · 19744c1 · 19744c1
1 parent ea0dcb7
commit 19744c1
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 9 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -16,7 +16,7 @@
  */
 
 /* tslint:disable */
-import {GoogleAuth} from 'google-auth-library';
+import {GoogleAuth, GoogleAuthOptions} from 'google-auth-library';
 
 import {
   processCountTokenResponse,
@@ -64,7 +64,8 @@ export class VertexAI {
     this.preview = new VertexAI_Preview(
       init.project,
       init.location,
-      init.apiEndpoint
+      init.apiEndpoint,
+      init.googleAuthOptions
     );
   }
 }
@@ -73,26 +74,48 @@ export class VertexAI {
  * VertexAI class internal implementation for authentication.
  */
 export class VertexAI_Preview {
-  protected googleAuth: GoogleAuth = new GoogleAuth({
-    scopes: 'https://www.googleapis.com/auth/cloud-platform',
-  });
+  protected googleAuth: GoogleAuth;
 
   /**
    * @constructor
    * @param {string} - project The Google Cloud project to use for the request
    * @param {string} - location The Google Cloud project location to use for the request
-   * @param {string} - apiEndpoint The base Vertex AI endpoint to use for the request. If
+   * @param {string} - [apiEndpoint] The base Vertex AI endpoint to use for the request. If
    *        not provided, the default regionalized endpoint
    *        (i.e. us-central1-aiplatform.googleapis.com) will be used.
+   * @param {GoogleAuthOptions} - [googleAuthOptions] The Authentication options provided by google-auth-library.
+   *        Complete list of authentication options is documented in the GoogleAuthOptions interface:
+   *        https://github.com/googleapis/google-auth-library-nodejs/blob/main/src/auth/googleauth.ts
    */
   constructor(
     readonly project: string,
     readonly location: string,
-    readonly apiEndpoint?: string
+    readonly apiEndpoint?: string,
+    readonly googleAuthOptions?: GoogleAuthOptions
   ) {
+    let opts: GoogleAuthOptions;
+    if (!googleAuthOptions) {
+      opts = {
+        scopes: 'https://www.googleapis.com/auth/cloud-platform',
+      };
+    } else {
+      if (
+        googleAuthOptions.projectId &&
+        googleAuthOptions.projectId !== project
+      ) {
+        throw new Error(
+          `inconsistent project ID values. argument project got value ${project} but googleAuthOptions.projectId got value ${googleAuthOptions.projectId}`
+        );
+      }
+      opts = googleAuthOptions;
+      if (!opts.scopes) {
+        opts.scopes = 'https://www.googleapis.com/auth/cloud-platform';
+      }
+    }
     this.project = project;
     this.location = location;
     this.apiEndpoint = apiEndpoint;
+    this.googleAuth = new GoogleAuth(opts);
   }
 
   /**

diff --git a/src/types/content.ts b/src/types/content.ts
@@ -15,16 +15,23 @@
  * limitations under the License.
  */
 
+// @ts-nocheck
+import {GoogleAuthOptions} from 'google-auth-library';
+
 /**
  * Params used to initialize the Vertex SDK
  * @param{string} project - the project name of your Google Cloud project. It is not the numeric project ID.
  * @param{string} location - the location of your project.
- * @param{string} apiEndpoint - Optional. If not specified, a default value will be resolved by SDK.
+ * @param{string} [apiEndpoint] - If not specified, a default value will be resolved by SDK.
+ * @param {GoogleAuthOptions} - [googleAuthOptions] The Authentication options provided by google-auth-library.
+ *        Complete list of authentication options is documented in the GoogleAuthOptions interface:
+ *        https://github.com/googleapis/google-auth-library-nodejs/blob/main/src/auth/googleauth.ts
  */
 export declare interface VertexInit {
   project: string;
   location: string;
   apiEndpoint?: string;
+  googleAuthOptions?: GoogleAuthOptions;
 }
 
 /**

diff --git a/test/index_test.ts b/test/index_test.ts
@@ -187,10 +187,39 @@ describe('VertexAI', () => {
     fetchSpy = spyOn(global, 'fetch').and.resolveTo(fetchResult);
   });
 
-  it('should be instantiated', () => {
+  it('given undefined google auth options, should be instantiated', () => {
     expect(vertexai).toBeInstanceOf(VertexAI);
   });
 
+  it('given specified google auth options, should be instantiated', () => {
+    const googleAuthOptions = {
+      scopes: 'test.scopes',
+    };
+    const vetexai1 = new VertexAI({
+      project: PROJECT,
+      location: LOCATION,
+      googleAuthOptions: googleAuthOptions,
+    });
+    expect(vetexai1).toBeInstanceOf(VertexAI);
+  });
+
+  it('given inconsistent project ID, should throw error', () => {
+    const googleAuthOptions = {
+      projectId: 'another_project',
+    };
+    expect(() => {
+      new VertexAI({
+        project: PROJECT,
+        location: LOCATION,
+        googleAuthOptions: googleAuthOptions,
+      });
+    }).toThrow(
+      new Error(
+        'inconsistent project ID values. argument project got value test_project but googleAuthOptions.projectId got value another_project'
+      )
+    );
+  });
+
   describe('generateContent', () => {
     it('returns a GenerateContentResponse', async () => {
       const req: GenerateContentRequest = {