Convert backend types to classes

Author: dlarocque

common/api-review/vertexai.api.md

@@ -12,6 +12,7 @@ import { FirebaseError } from '@firebase/util';
 // @public
 export interface AI {
     app: FirebaseApp;
+    // Warning: (ae-forgotten-export) The symbol "Backend" needs to be exported by the entry point index.d.ts
     backend: Backend;
     // @deprecated
     location: string;
@@ -74,9 +75,6 @@ export class ArraySchema extends Schema {
     toJSON(): SchemaRequest;
 }
 
-// @public
-export type Backend = GoogleAIBackend | VertexAIBackend;
-
 // @public
 export const BackendType: {
     readonly VERTEX_AI: "VERTEX_AI";
@@ -422,14 +420,6 @@ export function getImagenModel(ai: AI, modelParams: ImagenModelParams, requestOp
 // @public
 export function getVertexAI(app?: FirebaseApp, options?: VertexAIOptions): VertexAI;
 
-// @public
-export type GoogleAIBackend = {
-    backendType: typeof BackendType.GOOGLE_AI;
-};
-
-// @public
-export function googleAIBackend(): GoogleAIBackend;
-
 // @public @deprecated (undocumented)
 export interface GroundingAttribution {
     // (undocumented)
@@ -859,15 +849,6 @@ export interface UsageMetadata {
 // @public
 export type VertexAI = AI;
 
-// @public
-export type VertexAIBackend = {
-    backendType: typeof BackendType.VERTEX_AI;
-    location: string;
-};
-
-// @public
-export function vertexAIBackend(location?: string): VertexAIBackend;
-
 // @public
 export const VertexAIError: typeof AIError;
 

config/.eslintrc.js

@@ -174,7 +174,8 @@ module.exports = {
         }
       }
     ],
-    '@typescript-eslint/consistent-type-definitions': ['error', 'interface'],
+    // We prefer using interfaces, but we need to use types for aliases like '
+    // '@typescript-eslint/consistent-type-definitions': ['error', 'interface'],
     '@typescript-eslint/explicit-member-accessibility': [
       'error',
       {

packages/vertexai/src/api.test.ts

@@ -16,17 +16,11 @@
  */
 import { ImagenModelParams, ModelParams, AIErrorCode } from './types';
 import { AIError } from './errors';
-import {
-  ImagenModel,
-  getGenerativeModel,
-  getImagenModel,
-  googleAIBackend,
-  vertexAIBackend
-} from './api';
+import { ImagenModel, getGenerativeModel, getImagenModel } from './api';
 import { expect } from 'chai';
-import { BackendType, AI } from './public-types';
+import { AI } from './public-types';
 import { GenerativeModel } from './models/generative-model';
-import { DEFAULT_LOCATION } from './constants';
+import { VertexAIBackend } from './backend';
 
 const fakeAI: AI = {
   app: {
@@ -38,7 +32,7 @@ const fakeAI: AI = {
       appId: 'my-appid'
     }
   },
-  backend: vertexAIBackend('us-central1'),
+  backend: new VertexAIBackend('us-central1'),
   location: 'us-central1'
 };
 
@@ -171,28 +165,4 @@ describe('Top level API', () => {
     expect(genModel).to.be.an.instanceOf(ImagenModel);
     expect(genModel.model).to.equal('publishers/google/models/my-model');
   });
-  it('googleAIBackend returns a backend with backendType GOOGLE_AI', () => {
-    const backend = googleAIBackend();
-    expect(backend.backendType).to.equal(BackendType.GOOGLE_AI);
-  });
-  it('vertexAIBackend returns a backend with backendType VERTEX_AI', () => {
-    const backend = vertexAIBackend();
-    expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
-    expect(backend.location).to.equal(DEFAULT_LOCATION);
-  });
-  it('vertexAIBackend sets custom location', () => {
-    const backend = vertexAIBackend('test-location');
-    expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
-    expect(backend.location).to.equal('test-location');
-  });
-  it('vertexAIBackend sets custom location even if empty string', () => {
-    const backend = vertexAIBackend('');
-    expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
-    expect(backend.location).to.equal('');
-  });
-  it('vertexAIBackend uses default location if location is null', () => {
-    const backend = vertexAIBackend(null as any);
-    expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
-    expect(backend.location).to.equal(DEFAULT_LOCATION);
-  });
 });

packages/vertexai/src/api.ts

@@ -24,9 +24,7 @@ import {
   BackendType,
   AI,
   AIOptions,
-  GoogleAIBackend,
   VertexAI,
-  VertexAIBackend,
   VertexAIOptions
 } from './public-types';
 import {
@@ -38,6 +36,7 @@ import {
 import { AIError } from './errors';
 import { AIModel, GenerativeModel, ImagenModel } from './models';
 import { encodeInstanceIdentifier } from './helpers';
+import { GoogleAIBackend, VertexAIBackend } from './backend';
 
 export { ChatSession } from './methods/chat-session';
 export * from './requests/schema-builder';
@@ -72,7 +71,7 @@ declare module '@firebase/component' {
 
 /**
  * It is recommended to use the new {@link getAI | getAI()}.
- * 
+ *
  * Returns a {@link VertexAI} instance for the given app.
  *
  * @public
@@ -109,13 +108,13 @@ export function getVertexAI(
  * @example
  * ```javascript
  * // Get an AI instance configured to use Google AI.
- * const ai = getAI(app, { backend: googleAIBackend() });
+ * const ai = getAI(app, { backend: new GoogleAIBackend() });
  * ```
  *
  * @example
  * ```javascript
  * // Get an AI instance configured to use Vertex AI.
- * const ai = getAI(app, { backend: vertexAIBackend() });
+ * const ai = getAI(app, { backend: new VertexAIBackend() });
  * ```
  *
  * @param app - The {@link @firebase/app#FirebaseApp} to use.
@@ -126,52 +125,33 @@ export function getVertexAI(
  */
 export function getAI(
   app: FirebaseApp = getApp(),
-  options: AIOptions = { backend: googleAIBackend() }
+  options: AIOptions = { backend: new GoogleAIBackend() }
 ): AI {
   app = getModularInstance(app);
   // Dependencies
   const AIProvider: Provider<'AI'> = _getProvider(app, AI_TYPE);
 
-  const identifier = encodeInstanceIdentifier(options.backend);
+  let identifier: string;
+  if (options.backend instanceof GoogleAIBackend) {
+    identifier = encodeInstanceIdentifier({
+      backendType: BackendType.GOOGLE_AI
+    });
+  } else if (options.backend instanceof VertexAIBackend) {
+    identifier = encodeInstanceIdentifier({
+      backendType: BackendType.VERTEX_AI,
+      location: options.backend.location ?? DEFAULT_LOCATION
+    });
+  } else {
+    throw new AIError(
+      AIErrorCode.ERROR,
+      `Invalid backend type: ${options.backend.backendType}`
+    );
+  }
   return AIProvider.getImmediate({
     identifier
   });
 }
 
-/**
- * Creates a {@link Backend} instance configured to use Google AI.
- *
- * @returns A {@link GoogleAIBackend} object.
- *
- * @public
- */
-export function googleAIBackend(): GoogleAIBackend {
-  const backend: GoogleAIBackend = {
-    backendType: BackendType.GOOGLE_AI
-  };
-
-  return backend;
-}
-
-/**
- * Creates a {@link Backend} instance configured to use Vertex AI.
- *
- * @param location - The region identifier, defaulting to `us-central1`;
- * see {@link https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations | Vertex AI locations}
- * for a list of supported locations.
- * @returns A {@link VertexAIBackend} object.
- *
- * @public
- */
-export function vertexAIBackend(location?: string): VertexAIBackend {
-  const backend: VertexAIBackend = {
-    backendType: BackendType.VERTEX_AI,
-    location: location ?? DEFAULT_LOCATION
-  };
-
-  return backend;
-}
-
 /**
  * Returns a {@link GenerativeModel} class with methods for inference
  * and other functionality.

packages/vertexai/src/backend.test.ts

@@ -0,0 +1,35 @@
+import { expect } from "chai";
+import { GoogleAIBackend, VertexAIBackend } from "./backend";
+import { BackendType } from "./public-types";
+import { DEFAULT_LOCATION } from "./constants";
+
+describe('Backend', () => {
+  describe('GoogleAIBackend', () => {
+  it('sets backendType to GOOGLE_AI', () => {
+      const backend = new GoogleAIBackend();
+      expect(backend.backendType).to.equal(BackendType.GOOGLE_AI);
+    });
+  });
+  describe('VertexAIBackend', () => {
+    it('set backendType to VERTEX_AI', () => {
+      const backend = new VertexAIBackend();
+      expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
+      expect(backend.location).to.equal(DEFAULT_LOCATION);
+    });
+    it('sets custom location', () => {
+      const backend = new VertexAIBackend('test-location');
+      expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
+      expect(backend.location).to.equal('test-location');
+    });
+    it('sets custom location even if empty string', () => {
+      const backend = new VertexAIBackend('');
+      expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
+      expect(backend.location).to.equal('');
+    });
+    it('uses default location if location is null', () => {
+      const backend = new VertexAIBackend(null as any);
+      expect(backend.backendType).to.equal(BackendType.VERTEX_AI);
+      expect(backend.location).to.equal(DEFAULT_LOCATION);
+    });
+  });
+});

packages/vertexai/src/backend.ts

@@ -0,0 +1,72 @@
+import { DEFAULT_LOCATION } from "./constants";
+import { BackendType } from "./public-types";
+
+/**
+ * Abstract base class representing the configuration for an AI service backend.
+ * This class should not be instantiated directly. Use its subclasses
+ * {@link GoogleAIBackend} or {@link VertexAIBackend}.
+ *
+ * @public
+ */
+export abstract class Backend {
+  /**
+   * Specifies the backend type (either 'GOOGLE_AI' or 'VERTEX_AI').
+   */
+  readonly backendType: BackendType;
+
+  /**
+   * Protected constructor for use by subclasses.
+   * @param type - The specific backend type constant (e.g., BackendType.GOOGLE_AI).
+   */
+  protected constructor(type: BackendType) {
+    this.backendType = type;
+  }
+}
+
+/**
+ * Represents the configuration class for the Google AI backend.
+ * Use this with {@link AIOptions} when initializing the service with
+ * {@link getAI | getAI()}.
+ *
+ * @public
+ */
+export class GoogleAIBackend extends Backend {
+  /**
+   * Creates a configuration object for the Google AI backend.
+   */
+  constructor() {
+    super(BackendType.GOOGLE_AI);
+  }
+}
+
+/**
+ * Represents the configuration class for the Vertex AI backend.
+ * Use this with {@link AIOptions} when initializing the server with
+ * {@link getAI | getAI() }.
+ *
+ * @public
+ */
+export class VertexAIBackend extends Backend {
+  /**
+   * The region identifier.
+   * See {@link https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations | Vertex AI locations}
+   * for a list of supported locations.
+   */
+  readonly location: string;
+
+  /**
+   * Creates a configuration object for the Vertex AI backend.
+   *
+   * @param location - The region identifier, defaulting to `us-central1`;
+   * see {@link https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations | Vertex AI locations}
+   * for a list of supported locations.
+   */
+  constructor(location: string = DEFAULT_LOCATION) {
+    super(BackendType.VERTEX_AI);
+    if (location === null) {
+      this.location = DEFAULT_LOCATION;
+    } else {
+      this.location = location;
+    }
+  }
+}

packages/vertexai/src/backwards-compatbility.test.ts

@@ -24,10 +24,10 @@ import {
   VertexAIErrorCode,
   VertexAIModel,
   getGenerativeModel,
-  getImagenModel,
-  vertexAIBackend
+  getImagenModel
 } from './api';
 import { AI, VertexAI, AIErrorCode } from './public-types';
+import { VertexAIBackend } from './backend';
 
 function assertAssignable<T, _U extends T>(): void {}
 
@@ -41,7 +41,7 @@ const fakeAI: AI = {
       appId: 'app-id'
     }
   },
-  backend: vertexAIBackend('us-central1'),
+  backend: new VertexAIBackend('us-central1'),
   location: 'us-central1'
 };
 

packages/vertexai/src/errors.ts

@@ -16,7 +16,7 @@
  */
 
 import { FirebaseError } from '@firebase/util';
-import { AIErrorCode as AIErrorCode, CustomErrorData } from './types';
+import { AIErrorCode, CustomErrorData } from './types';
 import { VERTEX_TYPE } from './constants';
 
 /**

packages/vertexai/src/methods/chat-session.test.ts

@@ -23,7 +23,7 @@ import * as generateContentMethods from './generate-content';
 import { GenerateContentStreamResult } from '../types';
 import { ChatSession } from './chat-session';
 import { ApiSettings } from '../types/internal';
-import { vertexAIBackend } from '../api';
+import { VertexAIBackend } from '../backend';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -33,7 +33,7 @@ const fakeApiSettings: ApiSettings = {
   project: 'my-project',
   appId: 'my-appid',
   location: 'us-central1',
-  backend: vertexAIBackend()
+  backend: new VertexAIBackend()
 };
 
 describe('ChatSession', () => {