feat: infer location if user doesn't specifies it.

PiperOrigin-RevId: 635935845

Author: yyyu-google

src/types/content.ts

@@ -24,8 +24,12 @@ import {GoogleAuth, GoogleAuthOptions} from 'google-auth-library';
 export declare interface VertexInit {
   /** The Google Cloud project name. It is not the numeric project ID. */
   project: string;
-  /** The Google Cloud project location. */
-  location: string;
+  /**
+   * Optional. The Google Cloud project location. If not provided, SDK will
+   * firtly try to resolve it from run time environment. If no location resolved
+   * from run time environment, SDK will use default value `us-central1`.
+   */
+  location?: string;
   /**
    * Optional. The base Vertex AI endpoint to use for the request. If not
    * provided, the default regionalized endpoint (i.e.

src/vertex_ai.ts

@@ -52,7 +52,7 @@ export class VertexAI {
       init.googleAuthOptions
     );
     this.project = init.project;
-    this.location = init.location;
+    this.location = resolveLocation(init.location);
     this.googleAuth = new GoogleAuth(opts);
     this.apiEndpoint = init.apiEndpoint;
     this.preview = new VertexAIPreview(
@@ -222,3 +222,15 @@ function validateGoogleAuthOptions(
   }
   return opts;
 }
+
+function resolveLocation(locationFromInput?: string): string {
+  if (locationFromInput) {
+    return locationFromInput;
+  }
+  const inferredLocation =
+    process.env['GOOGLE_CLOUD_REGION'] || process.env['CLOUD_ML_REGION'];
+  if (inferredLocation) {
+    return inferredLocation;
+  }
+  return 'us-central1';
+}

system_test/end_to_end_sample_test.ts

@@ -27,7 +27,7 @@ import {
 } from '../src';
 import {FunctionDeclarationSchemaType} from '../src/types';
 
-const PROJECT = process.env.GCLOUD_PROJECT;
+const PROJECT = process.env['GCLOUD_PROJECT'];
 const LOCATION = 'us-central1';
 const TEXT_REQUEST = {
   contents: [{role: 'user', parts: [{text: 'How are you doing today?'}]}],

test/vertex_ai_test.ts

@@ -14,6 +14,50 @@ describe('VertexAI', () => {
     });
   });
 
+  it('no location given, should instantiate VertexAI and VertexAIPreview', () => {
+    const vertexaiNoLocation = new VertexAI({project: PROJECT});
+    const generativeModel = vertexaiNoLocation.getGenerativeModel({
+      model: 'gemini-pro',
+    });
+    const generativeModelPreview =
+      vertexaiNoLocation.preview.getGenerativeModel({
+        model: 'gemini-pro',
+      });
+    expect(vertexaiNoLocation).toBeInstanceOf(VertexAI);
+    expect(generativeModel).toBeInstanceOf(GenerativeModel);
+    expect(generativeModelPreview).toBeInstanceOf(GenerativeModelPreview);
+  });
+
+  it('location in run time env GOOGLE_CLOUD_REGION, should instantiate VertexAI and VertexAIPreview', () => {
+    process.env['GOOGLE_CLOUD_REGION'] = 'us-central1';
+    const vertexaiNoLocation = new VertexAI({project: PROJECT});
+    const generativeModel = vertexaiNoLocation.getGenerativeModel({
+      model: 'gemini-pro',
+    });
+    const generativeModelPreview =
+      vertexaiNoLocation.preview.getGenerativeModel({
+        model: 'gemini-pro',
+      });
+    expect(vertexaiNoLocation).toBeInstanceOf(VertexAI);
+    expect(generativeModel).toBeInstanceOf(GenerativeModel);
+    expect(generativeModelPreview).toBeInstanceOf(GenerativeModelPreview);
+  });
+
+  it('location in run time env CLOUD_ML_REGION, should instantiate VertexAI and VertexAIPreview', () => {
+    process.env['CLOUD_ML_REGION'] = 'us-central1';
+    const vertexaiNoLocation = new VertexAI({project: PROJECT});
+    const generativeModel = vertexaiNoLocation.getGenerativeModel({
+      model: 'gemini-pro',
+    });
+    const generativeModelPreview =
+      vertexaiNoLocation.preview.getGenerativeModel({
+        model: 'gemini-pro',
+      });
+    expect(vertexaiNoLocation).toBeInstanceOf(VertexAI);
+    expect(generativeModel).toBeInstanceOf(GenerativeModel);
+    expect(generativeModelPreview).toBeInstanceOf(GenerativeModelPreview);
+  });
+
   it('given undefined google auth options, should be instantiated', () => {
     expect(vertexai).toBeInstanceOf(VertexAI);
   });