Encode/decode instance identifiers directly to/from backends

common/api-review/vertexai.api.md

@@ -12,7 +12,6 @@ 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;
@@ -75,6 +74,12 @@ export class ArraySchema extends Schema {
     toJSON(): SchemaRequest;
 }
 
+// @public
+export abstract class Backend {
+    protected constructor(type: BackendType);
+    readonly backendType: BackendType;
+}
+
 // @public
 export const BackendType: {
     readonly VERTEX_AI: "VERTEX_AI";
@@ -418,6 +423,11 @@ export function getImagenModel(ai: AI, modelParams: ImagenModelParams, requestOp
 // @public
 export function getVertexAI(app?: FirebaseApp, options?: VertexAIOptions): VertexAI;
 
+// @public
+export class GoogleAIBackend extends Backend {
+    constructor();
+}
+
 // Warning: (ae-internal-missing-underscore) The name "GoogleAICitationMetadata" should be prefixed with an underscore because the declaration is marked as @internal
 //
 // @internal (undocumented)
@@ -897,6 +907,12 @@ export interface UsageMetadata {
 // @public
 export type VertexAI = AI;
 
+// @public
+export class VertexAIBackend extends Backend {
+    constructor(location?: string);
+    readonly location: string;
+}
+
 // @public
 export const VertexAIError: typeof AIError;
 

packages/vertexai/src/api.ts

@@ -18,15 +18,9 @@
 import { FirebaseApp, getApp, _getProvider } from '@firebase/app';
 import { Provider } from '@firebase/component';
 import { getModularInstance } from '@firebase/util';
-import { DEFAULT_LOCATION, AI_TYPE } from './constants';
+import { AI_TYPE } from './constants';
 import { AIService } from './service';
-import {
-  BackendType,
-  AI,
-  AIOptions,
-  VertexAI,
-  VertexAIOptions
-} from './public-types';
+import { AI, AIOptions, VertexAI, VertexAIOptions } from './public-types';
 import {
   ImagenModelParams,
   ModelParams,
@@ -42,6 +36,7 @@ export { ChatSession } from './methods/chat-session';
 export * from './requests/schema-builder';
 export { ImagenImageFormat } from './requests/imagen-image-format';
 export { AIModel, GenerativeModel, ImagenModel, AIError };
+export { Backend, VertexAIBackend, GoogleAIBackend } from './backend';
 
 export { AIErrorCode as VertexAIErrorCode };
 
@@ -86,10 +81,8 @@ export function getVertexAI(
   // Dependencies
   const AIProvider: Provider<'AI'> = _getProvider(app, AI_TYPE);
 
-  const identifier = encodeInstanceIdentifier({
-    backendType: BackendType.VERTEX_AI,
-    location: options?.location ?? DEFAULT_LOCATION
-  });
+  const backend = new VertexAIBackend(options?.location);
+  const identifier = encodeInstanceIdentifier(backend);
   return AIProvider.getImmediate({
     identifier
   });
@@ -131,22 +124,7 @@ export function getAI(
   // Dependencies
   const AIProvider: Provider<'AI'> = _getProvider(app, AI_TYPE);
 
-  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}`
-    );
-  }
+  const identifier = encodeInstanceIdentifier(options.backend);
   return AIProvider.getImmediate({
     identifier
   });

packages/vertexai/src/constants.ts

@@ -16,18 +16,12 @@
  */
 
 import { version } from '../package.json';
-import { BackendType } from './public-types';
-import { InstanceIdentifier } from './types/internal';
 
 // TODO (v12): Remove this
 export const VERTEX_TYPE = 'vertexAI';
 
 export const AI_TYPE = 'AI';
 
-export const DEFAULT_INSTANCE_IDENTIFIER: InstanceIdentifier = {
-  backendType: BackendType.GOOGLE_AI
-};
-
 export const DEFAULT_LOCATION = 'us-central1';
 
 export const DEFAULT_BASE_URL = 'https://firebasevertexai.googleapis.com';

packages/vertexai/src/helpers.test.ts

@@ -18,52 +18,39 @@ import { expect } from 'chai';
 import { AI_TYPE } from './constants';
 import { encodeInstanceIdentifier, decodeInstanceIdentifier } from './helpers';
 import { AIError } from './errors';
-import { BackendType } from './public-types';
-import { InstanceIdentifier } from './types/internal';
 import { AIErrorCode } from './types';
+import { GoogleAIBackend, VertexAIBackend } from './backend';
 
 describe('Identifier Encoding/Decoding', () => {
   describe('encodeInstanceIdentifier', () => {
     it('should encode Vertex AI identifier with a specific location', () => {
-      const identifier: InstanceIdentifier = {
-        backendType: BackendType.VERTEX_AI,
-        location: 'us-central1'
-      };
+      const backend = new VertexAIBackend('us-central1');
       const expected = `${AI_TYPE}/vertexai/us-central1`;
-      expect(encodeInstanceIdentifier(identifier)).to.equal(expected);
+      expect(encodeInstanceIdentifier(backend)).to.equal(expected);
     });
 
     it('should encode Vertex AI identifier using empty location', () => {
-      const identifier: InstanceIdentifier = {
-        backendType: BackendType.VERTEX_AI,
-        location: ''
-      };
+      const backend = new VertexAIBackend('');
       const expected = `${AI_TYPE}/vertexai/`;
-      expect(encodeInstanceIdentifier(identifier)).to.equal(expected);
+      expect(encodeInstanceIdentifier(backend)).to.equal(expected);
     });
 
     it('should encode Google AI identifier', () => {
-      const identifier: InstanceIdentifier = {
-        backendType: BackendType.GOOGLE_AI
-      };
+      const backend = new GoogleAIBackend();
       const expected = `${AI_TYPE}/googleai`;
-      expect(encodeInstanceIdentifier(identifier)).to.equal(expected);
+      expect(encodeInstanceIdentifier(backend)).to.equal(expected);
     });
 
     it('should throw AIError for unknown backend type', () => {
-      const identifier = {
-        backendType: 'some-future-backend'
-      } as any; // bypass type checking for the test
-
-      expect(() => encodeInstanceIdentifier(identifier)).to.throw(AIError);
+      expect(() => encodeInstanceIdentifier({} as any)).to.throw(AIError);
 
       try {
-        encodeInstanceIdentifier(identifier);
+        encodeInstanceIdentifier({} as any);
         expect.fail('Expected encodeInstanceIdentifier to throw');
       } catch (e) {
         expect(e).to.be.instanceOf(AIError);
         const error = e as AIError;
-        expect(error.message).to.contain(`Unknown backend`);
+        expect(error.message).to.contain('Invalid backend');
         expect(error.code).to.equal(AIErrorCode.ERROR);
       }
     });
@@ -72,11 +59,8 @@ describe('Identifier Encoding/Decoding', () => {
   describe('decodeInstanceIdentifier', () => {
     it('should decode Vertex AI identifier with location', () => {
       const encoded = `${AI_TYPE}/vertexai/europe-west1`;
-      const expected: InstanceIdentifier = {
-        backendType: BackendType.VERTEX_AI,
-        location: 'europe-west1'
-      };
-      expect(decodeInstanceIdentifier(encoded)).to.deep.equal(expected);
+      const backend = new VertexAIBackend('europe-west1');
+      expect(decodeInstanceIdentifier(encoded)).to.deep.equal(backend);
     });
 
     it('should throw an error if Vertex AI identifier string without explicit location part', () => {
@@ -98,10 +82,8 @@ describe('Identifier Encoding/Decoding', () => {
 
     it('should decode Google AI identifier', () => {
       const encoded = `${AI_TYPE}/googleai`;
-      const expected: InstanceIdentifier = {
-        backendType: BackendType.GOOGLE_AI
-      };
-      expect(decodeInstanceIdentifier(encoded)).to.deep.equal(expected);
+      const backend = new GoogleAIBackend();
+      expect(decodeInstanceIdentifier(encoded)).to.deep.equal(backend);
     });
 
     it('should throw AIError for invalid backend string', () => {

packages/vertexai/src/helpers.ts

@@ -17,50 +17,43 @@
 
 import { AI_TYPE } from './constants';
 import { AIError } from './errors';
-import { BackendType } from './public-types';
-import { InstanceIdentifier } from './types/internal';
 import { AIErrorCode } from './types';
+import { Backend, GoogleAIBackend, VertexAIBackend } from './backend';
 
 /**
- * Encodes an {@link InstanceIdentifier} into a string.
- *
- * This string is used to identify unique {@link AI} instances by backend type.
+ * Encodes a {@link Backend} into a string that will be used to uniquely identify {@link AI}
+ * instances by backend type.
  *
  * @internal
  */
-export function encodeInstanceIdentifier(
-  instanceIdentifier: InstanceIdentifier
-): string {
-  switch (instanceIdentifier.backendType) {
-    case BackendType.VERTEX_AI:
-      return `${AI_TYPE}/vertexai/${instanceIdentifier.location}`;
-    case BackendType.GOOGLE_AI:
-      return `${AI_TYPE}/googleai`;
-    default:
-      throw new AIError(
-        AIErrorCode.ERROR,
-        `Unknown backend '${instanceIdentifier}'`
-      );
+export function encodeInstanceIdentifier(backend: Backend): string {
+  if (backend instanceof GoogleAIBackend) {
+    return `${AI_TYPE}/googleai`;
+  } else if (backend instanceof VertexAIBackend) {
+    return `${AI_TYPE}/vertexai/${backend.location}`;
+  } else {
+    throw new AIError(
+      AIErrorCode.ERROR,
+      `Invalid backend: ${JSON.stringify(backend.backendType)}`
+    );
   }
 }
 
 /**
- * Decodes an instance identifier string into an {@link InstanceIdentifier}.
+ * Decodes an instance identifier string into a {@link Backend}.
  *
  * @internal
  */
-export function decodeInstanceIdentifier(
-  instanceIdentifier: string
-): InstanceIdentifier {
+export function decodeInstanceIdentifier(instanceIdentifier: string): Backend {
   const identifierParts = instanceIdentifier.split('/');
   if (identifierParts[0] !== AI_TYPE) {
     throw new AIError(
       AIErrorCode.ERROR,
       `Invalid instance identifier, unknown prefix '${identifierParts[0]}'`
     );
   }
-  const backend = identifierParts[1];
-  switch (backend) {
+  const backendType = identifierParts[1];
+  switch (backendType) {
     case 'vertexai':
       const location: string | undefined = identifierParts[2];
       if (!location) {
@@ -69,14 +62,9 @@ export function decodeInstanceIdentifier(
           `Invalid instance identifier, unknown location '${instanceIdentifier}'`
         );
       }
-      return {
-        backendType: BackendType.VERTEX_AI,
-        location
-      };
+      return new VertexAIBackend(location);
     case 'googleai':
-      return {
-        backendType: BackendType.GOOGLE_AI
-      };
+      return new GoogleAIBackend();
     default:
       throw new AIError(
         AIErrorCode.ERROR,

packages/vertexai/src/index.node.ts

@@ -23,33 +23,31 @@
 
 import { registerVersion, _registerComponent } from '@firebase/app';
 import { AIService } from './service';
-import { DEFAULT_INSTANCE_IDENTIFIER, AI_TYPE } from './constants';
+import { AI_TYPE } from './constants';
 import { Component, ComponentType } from '@firebase/component';
 import { name, version } from '../package.json';
-import { InstanceIdentifier } from './types/internal';
 import { decodeInstanceIdentifier } from './helpers';
+import { AIError } from './errors';
+import { AIErrorCode } from './public-types';
 
 function registerAI(): void {
   _registerComponent(
     new Component(
       AI_TYPE,
-      (container, options) => {
-        // getImmediate for FirebaseApp will always succeed
-        const app = container.getProvider('app').getImmediate();
-        const auth = container.getProvider('auth-internal');
-        const appCheckProvider = container.getProvider('app-check-internal');
-
-        let instanceIdentifier: InstanceIdentifier;
-        if (options.instanceIdentifier) {
-          instanceIdentifier = decodeInstanceIdentifier(
-            options.instanceIdentifier
+      (container, { instanceIdentifier }) => {
+        if (!instanceIdentifier) {
+          throw new AIError(
+            AIErrorCode.ERROR,
+            'AIService instance identifier is undefined.'
           );
-        } else {
-          instanceIdentifier = DEFAULT_INSTANCE_IDENTIFIER;
         }
 
-        const backend = instanceIdentifier;
+        const backend = decodeInstanceIdentifier(instanceIdentifier);
 
+        // getImmediate for FirebaseApp will always succeed
+        const app = container.getProvider('app').getImmediate();
+        const auth = container.getProvider('auth-internal');
+        const appCheckProvider = container.getProvider('app-check-internal');
         return new AIService(app, backend, auth, appCheckProvider);
       },
       ComponentType.PUBLIC

packages/vertexai/src/index.ts

@@ -49,6 +49,7 @@ function registerAI(): void {
         }
 
         const backend = decodeInstanceIdentifier(instanceIdentifier);
+
         // getImmediate for FirebaseApp will always succeed
         const app = container.getProvider('app').getImmediate();
         const auth = container.getProvider('auth-internal');

packages/vertexai/src/requests/request.ts

@@ -66,7 +66,7 @@ export class RequestUrl {
     } else {
       throw new AIError(
         AIErrorCode.ERROR,
-        `Invalid backend: ${this.apiSettings.backend}`
+        `Invalid backend: ${JSON.stringify(this.apiSettings.backend)}`
       );
     }
   }

packages/vertexai/src/types/internal.ts

@@ -18,7 +18,6 @@
 import { AppCheckTokenResult } from '@firebase/app-check-interop-types';
 import { FirebaseAuthTokenData } from '@firebase/auth-interop-types';
 import { Backend } from '../backend';
-import { BackendType } from '../public-types';
 
 export * from './imagen/internal';
 
@@ -28,15 +27,10 @@ export interface ApiSettings {
   appId: string;
   automaticDataCollectionEnabled?: boolean;
   /**
-   * @deprecated
+   * @deprecated Use `backend.location` instead.
    */
   location: string;
   backend: Backend;
   getAuthToken?: () => Promise<FirebaseAuthTokenData | null>;
   getAppCheckToken?: () => Promise<AppCheckTokenResult>;
 }
-
-export interface InstanceIdentifier {
-  backendType: BackendType;
-  location?: string;
-}