feat: allow requests to generativelanguage endpoint

This commit allows users to opt-in to using the  endpoint for Vertex AI requests, instead of the default  endpoint.

This change involved:
- Adding a new  configuration option to .
- Adding  to the initialization of
- Modifying  to use the new option.
- Adding the  and  functions.
- Adding the  and  functions.
- Adding new types for the  API.
- Adding unit tests to verify the new functionality.

Author: google-labs-jules[bot]

packages/vertexai/src/api.ts

@@ -58,7 +58,8 @@ export function getVertexAI(
   const vertexProvider: Provider<'vertexAI'> = _getProvider(app, VERTEX_TYPE);
 
   return vertexProvider.getImmediate({
-    identifier: options?.location || DEFAULT_LOCATION
+    identifier: options?.location || DEFAULT_LOCATION,
+    optionalArguments: options
   });
 }
 

packages/vertexai/src/requests/request.ts

@@ -45,7 +45,7 @@ export class RequestUrl {
   toString(): string {
     // TODO: allow user-set option if that feature becomes available
     const apiVersion = DEFAULT_API_VERSION;
-    const baseUrl = this.requestOptions?.baseUrl || DEFAULT_BASE_URL;
+    const baseUrl = this.requestOptions?.baseUrl || this.requestOptions?.vertexAiEndpoint || DEFAULT_BASE_URL;
     let url = `${baseUrl}/${apiVersion}`;
     url += `/projects/${this.apiSettings.project}`;
     url += `/locations/${this.apiSettings.location}`;
@@ -137,12 +137,16 @@ export async function makeRequest(
   let response;
   let fetchTimeoutId: string | number | NodeJS.Timeout | undefined;
   try {
+    let translatedBody = body;
+    if(requestOptions?.vertexAiEndpoint === 'generativelanguage.googleapis.com'){
+      translatedBody = JSON.stringify(translateRequest(JSON.parse(body)));
+    }
     const request = await constructRequest(
       model,
       task,
       apiSettings,
       stream,
-      body,
+      translatedBody,
       requestOptions
     );
     // Timeout is 180s by default
@@ -158,8 +162,12 @@ export async function makeRequest(
     if (!response.ok) {
       let message = '';
       let errorDetails;
+      let translatedResponse = await response.json()
+      if(requestOptions?.vertexAiEndpoint === 'generativelanguage.googleapis.com'){
+        translatedResponse = translateResponse(translatedResponse)
+      }
       try {
-        const json = await response.json();
+        const json = translatedResponse.error ? translatedResponse : {candidates: [translatedResponse]};
         message = json.error.message;
         if (json.error.details) {
           message += ` ${JSON.stringify(json.error.details)}`;

packages/vertexai/src/service.ts

@@ -32,6 +32,7 @@ export class VertexAIService implements VertexAI, _FirebaseService {
   auth: FirebaseAuthInternal | null;
   appCheck: FirebaseAppCheckInternal | null;
   location: string;
+  vertexAiEndpoint?: string;
 
   constructor(
     public app: FirebaseApp,
@@ -44,6 +45,7 @@ export class VertexAIService implements VertexAI, _FirebaseService {
     this.auth = auth || null;
     this.appCheck = appCheck || null;
     this.location = this.options?.location || DEFAULT_LOCATION;
+    this.vertexAiEndpoint = this.options?.vertexAiEndpoint;
   }
 
   _delete(): Promise<void> {

packages/vertexai/src/translation/index.ts

@@ -0,0 +1,145 @@
+/**
+ * @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 { Content, Part } from '../types/content';
+
+/**
+ * This is a function that translates a request to the
+ * `firebasevertexai` format into a request to the `generativelanguage`
+ * format.
+ */
+import { Content, Part } from '../types/content';
+import * as GenerativeLanguage from '../types/generativelanguage';
+
+/**
+ * This is a function that translates a request to the
+ * `firebasevertexai` format into a request to the `generativelanguage`
+ * format.
+ */
+export function translateRequest(
+  request: any
+): GenerativeLanguage.GenerateContentRequest {
+  const translatedContents = request.contents.map(translateContent);
+  return {
+    contents: translatedContents
+  };
+}
+
+/**
+ * This is a function that translates a response from the
+ * `generativelanguage` format into a response in the
+ * `firebasevertexai` format.
+ */
+import { Content, Part } from '../types/content';
+import * as GenerativeLanguage from '../types/generativelanguage';
+import { GenerateContentResponse } from '../types';
+
+/**
+ * This is a function that translates a request to the
+ * `firebasevertexai` format into a request to the `generativelanguage`
+ * format.
+ */
+export function translateRequest(
+  request: any
+): GenerativeLanguage.GenerateContentRequest {
+  const translatedContents = request.contents.map(translateContent);
+  return {
+    contents: translatedContents
+  };
+}
+
+/**
+ * This is a function that translates a response from the
+ * `generativelanguage` format into a response in the
+ * `firebasevertexai` format.
+ */
+export function translateResponse(
+  response: any
+): GenerateContentResponse {
+  const translatedResponse = {
+    candidates: [{
+        content: {
+            role: response[0].candidates[0].content.role,
+            parts: response[0].candidates[0].content.parts.map(translatePart)
+        },
+        finishReason: response[0].candidates[0].finishReason,
+        safetyRatings: response[0].candidates[0].safetyRatings,
+      }]
+  };
+
+  return translatedResponse;
+}
+
+/**
+ * This is a function that translates a content to the
+ * `generativelanguage` format into a content in the
+ * `firebasevertexai` format.
+ */
+export function translateContent(
+  content: Content | Part
+): GenerativeLanguage.Content {
+    const newContent = {
+      role: content.role,
+      parts: content.parts.map(translatePart)
+    }
+    return newContent;
+}
+
+/**
+ * This is a function that translates a part to the
+ * `generativelanguage` format into a part in the
+ * `firebasevertexai` format.
+ */
+export function translatePart(
+  part: Part
+): GenerativeLanguage.Part {
+  const newPart = {
+    text: part.text,
+    inlineData: part.inlineData
+  }
+  return newPart;
+}
+
+/**
+ * This is a function that translates a content to the
+ * `generativelanguage` format into a content in the
+ * `firebasevertexai` format.
+ */
+export function translateContent(
+  content: Content | Part
+): GenerativeLanguage.Content {
+    const newContent = {
+      role: content.role,
+      parts: content.parts.map(translatePart)
+    }
+    return newContent;
+}
+
+/**
+ * This is a function that translates a part to the
+ * `generativelanguage` format into a part in the
+ * `firebasevertexai` format.
+ */
+export function translatePart(
+  part: Part
+): GenerativeLanguage.Part {
+  const newPart = {
+    text: part.text,
+    inlineData: part.inlineData
+  }
+  return newPart;
+}

packages/vertexai/src/types/generativelanguage.ts

@@ -0,0 +1,33 @@
+/**
+ * @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.
+ */
+
+export interface GenerateContentRequest {
+  contents: Content[];
+}
+
+export interface Content {
+  role: string;
+  parts: Part[];
+}
+
+export interface Part {
+  text?: string;
+  inlineData?: {
+    mimeType: string;
+    data: string;
+  };
+}

packages/vertexai/src/types/requests.ts

@@ -129,6 +129,13 @@ export interface RequestOptions {
    * Base url for endpoint. Defaults to https://firebasevertexai.googleapis.com
    */
   baseUrl?: string;
+  /**
+   * If set, will use the new `generativelanguage` endpoint instead of `firebasevertexai`
+   * Should be one of `firebasevertexai.googleapis.com` or `generativelanguage.googleapis.com`
+   *
+   * @remarks This setting is used to change the Vertex AI API endpoint.
+   */
+  vertexAiEndpoint?: string;
 }
 
 /**

packages/vertexai/test/translation.test.ts

@@ -0,0 +1,96 @@
+/**
+ * @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 { translateRequest, translateResponse, translateContent, translatePart } from '../src/translation';
+import { GenerativeModel } from '../src/models';
+import { VertexAI } from '../src/public-types';
+import { RequestOptions } from '../src/types';
+import { getVertexAI } from '../src/api';
+import { FirebaseApp } from '@firebase/app';
+import { initializeApp } from '@firebase/app';
+import * as assert from 'assert';
+
+describe('VertexAI translation tests', () => {
+    let app: FirebaseApp;
+    let vertex: VertexAI;
+    let options: RequestOptions;
+    let model: GenerativeModel;
+    beforeEach(() => {
+        app = initializeApp({projectId:'test-project'});
+        vertex = getVertexAI(app, { location: 'us-central1' });
+        options = {vertexAiEndpoint : 'generativelanguage.googleapis.com', baseUrl:'https://test.googleapis.com/'}
+        model = new GenerativeModel(vertex, {model:'gemini-pro'}, options);
+    });
+    it('should use generativelanguage as endpoint', async () => {
+        const request = {
+          contents: [{
+              role: 'user',
+              parts: [{ text: 'hello' }]
+            }]
+        };
+        const translated = translateRequest(request);
+        const response = await model.generateContent(translated);
+
+        const expectedResponse = {
+            candidates:[{
+                content: {
+                    role: 'user',
+                    parts: [{ text: 'hello' }]
+                },
+                finishReason: 'STOP'
+            }],
+        };
+
+      it('should use the default url if no endpoint is specified', async () => {
+        const requestOptions: RequestOptions = {baseUrl:'https://test.googleapis.com/'};
+        const service = new VertexAIService(app);
+        let testurl = new URL("https://test.googleapis.com/");
+        let defaulturl = new URL("https://firebasevertexai.googleapis.com/");
+        const url1 = new RequestUrl('gemini-pro', 'generateContent', {project: 'test-project', location:'us-central1', apiKey:'test-key'}, false, requestOptions);
+        assert.strictEqual(url1.toString(), testurl.toString());
+        const url2 = new RequestUrl('gemini-pro', 'generateContent', {project: 'test-project', location:'us-central1', apiKey:'test-key'}, false, {});
+      it('should use the default url if no endpoint is specified', async () => {
+        const requestOptions: RequestOptions = {baseUrl:'https://test.googleapis.com/'};
+        const service = new VertexAIService(app);
+        let testurl = new URL("https://test.googleapis.com/");
+        let defaulturl = new URL("https://firebasevertexai.googleapis.com/");
+        const url1 = new RequestUrl('gemini-pro', 'generateContent', {project: 'test-project', location:'us-central1', apiKey:'test-key'}, false, requestOptions);
+        assert.strictEqual(url1.toString(), testurl.toString());
+        const url2 = new RequestUrl('gemini-pro', 'generateContent', {project: 'test-project', location:'us-central1', apiKey:'test-key'}, false, {});
+        assert.strictEqual(url2.toString(), defaulturl.toString());
+      });
+      it('should translate response correctly', () => {
+        const testResponse = {
+          response: {
+            candidates: [
+              {
+                content: {
+                  role: 'user',
+                  parts: [{ text: 'hello' }],
+                },
+                safetyRatings: [],
+                finishReason: 'STOP',
+              },
+            ]
+          }
+        };
+        const translated = translateResponse(testResponse);
+        assert.strictEqual(translated.candidates[0].content.role, 'user');
+        assert.strictEqual(translated.candidates[0].content.parts[0].text, 'hello');
+        assert.strictEqual(translated.candidates[0].finishReason, 'STOP');
+      });
+});