[Vertex AI] Add apiVersion parameter to RequestOptions

Author: andrewheard

packages/vertexai/src/requests/request.test.ts

@@ -22,7 +22,7 @@ import chaiAsPromised from 'chai-as-promised';
 import { RequestUrl, Task, getHeaders, makeRequest } from './request';
 import { ApiSettings } from '../types/internal';
 import { DEFAULT_API_VERSION } from '../constants';
-import { VertexAIErrorCode } from '../types';
+import { ApiVersion, VertexAIErrorCode } from '../types';
 import { VertexAIError } from '../errors';
 import { getMockResponse } from '../../test-utils/mock-response';
 
@@ -74,6 +74,18 @@ describe('request methods', () => {
       );
       expect(url.toString()).to.include(DEFAULT_API_VERSION);
     });
+    it('custom apiVersion', async () => {
+      const url = new RequestUrl(
+        'models/model-name',
+        Task.GENERATE_CONTENT,
+        fakeApiSettings,
+        false,
+        { apiVersion: ApiVersion.V1 }
+      );
+      expect(url.toString()).to.include(ApiVersion.V1);
+      // TODO: Update test to set ApiVersion.V1BETA when ApiVersion.V1 becomes the default.
+      expect(ApiVersion.V1).to.not.equal(DEFAULT_API_VERSION);
+    });
     it('custom baseUrl', async () => {
       const url = new RequestUrl(
         'models/model-name',

packages/vertexai/src/requests/request.ts

@@ -42,8 +42,7 @@ export class RequestUrl {
     public requestOptions?: RequestOptions
   ) {}
   toString(): string {
-    // TODO: allow user-set option if that feature becomes available
-    const apiVersion = DEFAULT_API_VERSION;
+    const apiVersion = this.requestOptions?.apiVersion || DEFAULT_API_VERSION;
     const baseUrl = this.requestOptions?.baseUrl || DEFAULT_BASE_URL;
     let url = `${baseUrl}/${apiVersion}`;
     url += `/projects/${this.apiSettings.project}`;

packages/vertexai/src/types/enums.ts

@@ -27,6 +27,17 @@ export type Role = (typeof POSSIBLE_ROLES)[number];
  */
 export const POSSIBLE_ROLES = ['user', 'model', 'function', 'system'] as const;
 
+/**
+ * Supported API versions for the Vertex AI in Firebase endpoint.
+ * @public
+ */
+export enum ApiVersion {
+  /** The stable channel for version 1 of the API. */
+  V1 = 'v1',
+  /** The beta channel for version 1 of the API. */
+  V1BETA = 'v1beta'
+}
+
 /**
  * Harm categories that would cause prompts or candidates to be blocked.
  * @public

packages/vertexai/src/types/requests.ts

@@ -18,6 +18,7 @@
 import { TypedSchema } from '../requests/schema-builder';
 import { Content, Part } from './content';
 import {
+  ApiVersion,
   FunctionCallingMode,
   HarmBlockMethod,
   HarmBlockThreshold,
@@ -129,6 +130,11 @@ export interface RequestOptions {
    * Base url for endpoint. Defaults to https://firebasevertexai.googleapis.com
    */
   baseUrl?: string;
+  /**
+   * API version for endpoint. Defaults to <code>"v1beta"</code>
+   * (<code>{@link ApiVersion.V1BETA}</code>).
+   */
+  apiVersion?: ApiVersion;
 }
 
 /**