Merge a49dc31 into b981a45

Author: hsubox76

packages/vertexai/src/api.ts

@@ -20,7 +20,7 @@ import { Provider } from '@firebase/component';
 import { getModularInstance } from '@firebase/util';
 import { DEFAULT_LOCATION, VERTEX_TYPE } from './constants';
 import { VertexAIService } from './service';
-import { VertexAI } from './public-types';
+import { VertexAI, VertexAIOptions } from './public-types';
 import { ERROR_FACTORY, VertexError } from './errors';
 import { ModelParams, RequestOptions } from './types';
 import { GenerativeModel } from './models/generative-model';
@@ -42,13 +42,16 @@ declare module '@firebase/component' {
  *
  * @param app - The {@link @firebase/app#FirebaseApp} to use.
  */
-export function getVertexAI(app: FirebaseApp = getApp()): VertexAI {
+export function getVertexAI(
+  app: FirebaseApp = getApp(),
+  options?: VertexAIOptions
+): VertexAI {
   app = getModularInstance(app);
   // Dependencies
   const vertexProvider: Provider<'vertex'> = _getProvider(app, VERTEX_TYPE);
 
   return vertexProvider.getImmediate({
-    identifier: DEFAULT_LOCATION
+    identifier: options?.location || DEFAULT_LOCATION
   });
 }
 

packages/vertexai/src/index.ts

@@ -37,11 +37,11 @@ function registerVertex(): void {
   _registerComponent(
     new Component(
       VERTEX_TYPE,
-      container => {
+      (container, { instanceIdentifier: location }) => {
         // getImmediate for FirebaseApp will always succeed
         const app = container.getProvider('app').getImmediate();
         const appCheckProvider = container.getProvider('app-check-internal');
-        return new VertexAIService(app, appCheckProvider);
+        return new VertexAIService(app, appCheckProvider, { location });
       },
       ComponentType.PUBLIC
     ).setMultipleInstances(true)

packages/vertexai/src/public-types.ts

@@ -30,3 +30,11 @@ export interface VertexAI {
   app: FirebaseApp;
   location: string;
 }
+
+/**
+ * Options when initializing the Firebase Vertex AI SDK.
+ * @public
+ */
+export interface VertexAIOptions {
+  location?: string;
+}

packages/vertexai/src/requests/stream-reader.test.ts

@@ -32,7 +32,8 @@ import {
   FinishReason,
   GenerateContentResponse,
   HarmCategory,
-  HarmProbability
+  HarmProbability,
+  SafetyRating
 } from '../types';
 
 use(sinonChai);
@@ -229,7 +230,7 @@ describe('aggregateResponses', () => {
             {
               category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
               probability: HarmProbability.LOW
-            }
+            } as SafetyRating
           ]
         }
       }
@@ -256,7 +257,7 @@ describe('aggregateResponses', () => {
                 {
                   category: HarmCategory.HARM_CATEGORY_HARASSMENT,
                   probability: HarmProbability.NEGLIGIBLE
-                }
+                } as SafetyRating
               ]
             }
           ],
@@ -266,7 +267,7 @@ describe('aggregateResponses', () => {
               {
                 category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
                 probability: HarmProbability.LOW
-              }
+              } as SafetyRating
             ]
           }
         },
@@ -284,7 +285,7 @@ describe('aggregateResponses', () => {
                 {
                   category: HarmCategory.HARM_CATEGORY_HARASSMENT,
                   probability: HarmProbability.NEGLIGIBLE
-                }
+                } as SafetyRating
               ],
               citationMetadata: {
                 citations: [
@@ -304,7 +305,7 @@ describe('aggregateResponses', () => {
               {
                 category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
                 probability: HarmProbability.HIGH
-              }
+              } as SafetyRating
             ]
           }
         },
@@ -322,7 +323,7 @@ describe('aggregateResponses', () => {
                 {
                   category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
                   probability: HarmProbability.MEDIUM
-                }
+                } as SafetyRating
               ],
               citationMetadata: {
                 citations: [
@@ -348,7 +349,7 @@ describe('aggregateResponses', () => {
               {
                 category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
                 probability: HarmProbability.HIGH
-              }
+              } as SafetyRating
             ]
           }
         }

packages/vertexai/src/service.test.ts

@@ -0,0 +1,43 @@
+/**
+ * @license
+ * 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 { DEFAULT_LOCATION } from './constants';
+import { VertexAIService } from './service';
+import { expect } from 'chai';
+
+const fakeApp = {
+  name: 'DEFAULT',
+  automaticDataCollectionEnabled: true,
+  options: {
+    apiKey: 'key',
+    projectId: 'my-project'
+  }
+};
+
+describe('VertexAIService', () => {
+  it('uses default location if not specified', () => {
+    const vertexAI = new VertexAIService(fakeApp);
+    expect(vertexAI.location).to.equal(DEFAULT_LOCATION);
+  });
+  it('uses custom location if specified', () => {
+    const vertexAI = new VertexAIService(
+      fakeApp,
+      /* appCheckProvider */ undefined,
+      { location: 'somewhere' }
+    );
+    expect(vertexAI.location).to.equal('somewhere');
+  });
+});

packages/vertexai/src/service.ts

@@ -16,7 +16,7 @@
  */
 
 import { FirebaseApp, _FirebaseService } from '@firebase/app';
-import { VertexAI } from './public-types';
+import { VertexAI, VertexAIOptions } from './public-types';
 import {
   AppCheckInternalComponentName,
   FirebaseAppCheckInternal
@@ -30,12 +30,12 @@ export class VertexAIService implements VertexAI, _FirebaseService {
 
   constructor(
     public app: FirebaseApp,
-    appCheckProvider?: Provider<AppCheckInternalComponentName>
+    appCheckProvider?: Provider<AppCheckInternalComponentName>,
+    public options?: VertexAIOptions
   ) {
     const appCheck = appCheckProvider?.getImmediate({ optional: true });
     this.appCheck = appCheck || null;
-    // TODO: add in user-set location option when that feature is available
-    this.location = DEFAULT_LOCATION;
+    this.location = this.options?.location || DEFAULT_LOCATION;
   }
 
   _delete(): Promise<void> {