Add X-Firebase-AppId header to VertexAI requests

.changeset/red-hornets-peel.md

@@ -0,0 +1,5 @@
+---
+'@firebase/vertexai': patch
+---
+
+Throw an error when initializing models if `appId` is not defined in the given `VertexAI` instance.

common/api-review/vertexai.api.md

@@ -817,6 +817,7 @@ export const enum VertexAIErrorCode {
     INVALID_CONTENT = "invalid-content",
     INVALID_SCHEMA = "invalid-schema",
     NO_API_KEY = "no-api-key",
+    NO_APP_ID = "no-app-id",
     NO_MODEL = "no-model",
     NO_PROJECT_ID = "no-project-id",
     PARSE_FAILED = "parse-failed",

docs-devsite/vertexai.md

@@ -545,6 +545,7 @@ export declare const enum VertexAIErrorCode
 |  INVALID\_CONTENT | <code>&quot;invalid-content&quot;</code> | An error associated with a Content object. |
 |  INVALID\_SCHEMA | <code>&quot;invalid-schema&quot;</code> | An error due to invalid Schema input. |
 |  NO\_API\_KEY | <code>&quot;no-api-key&quot;</code> | An error occurred due to a missing Firebase API key. |
+|  NO\_APP\_ID | <code>&quot;no-app-id&quot;</code> | An error occured due to a missing app ID. |
 |  NO\_MODEL | <code>&quot;no-model&quot;</code> | An error occurred due to a model name not being specified during initialization. |
 |  NO\_PROJECT\_ID | <code>&quot;no-project-id&quot;</code> | An error occurred due to a missing project ID. |
 |  PARSE\_FAILED | <code>&quot;parse-failed&quot;</code> | An error occurred while parsing. |

packages/vertexai/src/api.test.ts

@@ -27,7 +27,8 @@ const fakeVertexAI: VertexAI = {
     automaticDataCollectionEnabled: true,
     options: {
       apiKey: 'key',
-      projectId: 'my-project'
+      projectId: 'my-project',
+      appId: 'my-appid'
     }
   },
   location: 'us-central1'
@@ -48,7 +49,7 @@ describe('Top level API', () => {
   it('getGenerativeModel throws if no apiKey is provided', () => {
     const fakeVertexNoApiKey = {
       ...fakeVertexAI,
-      app: { options: { projectId: 'my-project' } }
+      app: { options: { projectId: 'my-project', appId: 'my-appid' } }
     } as VertexAI;
     try {
       getGenerativeModel(fakeVertexNoApiKey, { model: 'my-model' });
@@ -64,7 +65,7 @@ describe('Top level API', () => {
   it('getGenerativeModel throws if no projectId is provided', () => {
     const fakeVertexNoProject = {
       ...fakeVertexAI,
-      app: { options: { apiKey: 'my-key' } }
+      app: { options: { apiKey: 'my-key', appId: 'my-appid' } }
     } as VertexAI;
     try {
       getGenerativeModel(fakeVertexNoProject, { model: 'my-model' });
@@ -79,6 +80,22 @@ describe('Top level API', () => {
       );
     }
   });
+  it('getGenerativeModel throws if no appId is provided', () => {
+    const fakeVertexNoProject = {
+      ...fakeVertexAI,
+      app: { options: { apiKey: 'my-key', projectId: 'my-projectid' } }
+    } as VertexAI;
+    try {
+      getGenerativeModel(fakeVertexNoProject, { model: 'my-model' });
+    } catch (e) {
+      expect((e as VertexAIError).code).includes(VertexAIErrorCode.NO_APP_ID);
+      expect((e as VertexAIError).message).equals(
+        `VertexAI: The "appId" field is empty in the local` +
+          ` Firebase config. Firebase VertexAI requires this field ` +
+          `to contain a valid app ID. (vertexAI/${VertexAIErrorCode.NO_APP_ID})`
+      );
+    }
+  });
   it('getGenerativeModel gets a GenerativeModel', () => {
     const genModel = getGenerativeModel(fakeVertexAI, { model: 'my-model' });
     expect(genModel).to.be.an.instanceOf(GenerativeModel);
@@ -98,7 +115,7 @@ describe('Top level API', () => {
   it('getImagenModel throws if no apiKey is provided', () => {
     const fakeVertexNoApiKey = {
       ...fakeVertexAI,
-      app: { options: { projectId: 'my-project' } }
+      app: { options: { projectId: 'my-project', appId: 'my-appid' } }
     } as VertexAI;
     try {
       getImagenModel(fakeVertexNoApiKey, { model: 'my-model' });
@@ -114,7 +131,7 @@ describe('Top level API', () => {
   it('getImagenModel throws if no projectId is provided', () => {
     const fakeVertexNoProject = {
       ...fakeVertexAI,
-      app: { options: { apiKey: 'my-key' } }
+      app: { options: { apiKey: 'my-key', appId: 'my-appid' } }
     } as VertexAI;
     try {
       getImagenModel(fakeVertexNoProject, { model: 'my-model' });
@@ -129,6 +146,22 @@ describe('Top level API', () => {
       );
     }
   });
+  it('getImagenModel throws if no appId is provided', () => {
+    const fakeVertexNoProject = {
+      ...fakeVertexAI,
+      app: { options: { apiKey: 'my-key', projectId: 'my-project' } }
+    } as VertexAI;
+    try {
+      getImagenModel(fakeVertexNoProject, { model: 'my-model' });
+    } catch (e) {
+      expect((e as VertexAIError).code).includes(VertexAIErrorCode.NO_APP_ID);
+      expect((e as VertexAIError).message).equals(
+        `VertexAI: The "appId" field is empty in the local` +
+          ` Firebase config. Firebase VertexAI requires this field ` +
+          `to contain a valid app ID. (vertexAI/${VertexAIErrorCode.NO_APP_ID})`
+      );
+    }
+  });
   it('getImagenModel gets an ImagenModel', () => {
     const genModel = getImagenModel(fakeVertexAI, { model: 'my-model' });
     expect(genModel).to.be.an.instanceOf(ImagenModel);

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

@@ -30,6 +30,7 @@ use(chaiAsPromised);
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',
+  appId: 'my-appid',
   location: 'us-central1'
 };
 

packages/vertexai/src/methods/count-tokens.test.ts

@@ -32,6 +32,7 @@ use(chaiAsPromised);
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',
+  appId: 'my-appid',
   location: 'us-central1'
 };
 

packages/vertexai/src/methods/generate-content.test.ts

@@ -37,6 +37,7 @@ use(chaiAsPromised);
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',
+  appId: 'my-appid',
   location: 'us-central1'
 };
 

packages/vertexai/src/models/generative-model.test.ts

@@ -30,7 +30,8 @@ const fakeVertexAI: VertexAI = {
     automaticDataCollectionEnabled: true,
     options: {
       apiKey: 'key',
-      projectId: 'my-project'
+      projectId: 'my-project',
+      appId: 'my-appid'
     }
   },
   location: 'us-central1'

packages/vertexai/src/models/imagen-model.test.ts

@@ -37,7 +37,8 @@ const fakeVertexAI: VertexAI = {
     automaticDataCollectionEnabled: true,
     options: {
       apiKey: 'key',
-      projectId: 'my-project'
+      projectId: 'my-project',
+      appId: 'my-appid'
     }
   },
   location: 'us-central1'

packages/vertexai/src/models/vertexai-model.test.ts

@@ -38,7 +38,8 @@ const fakeVertexAI: VertexAI = {
     automaticDataCollectionEnabled: true,
     options: {
       apiKey: 'key',
-      projectId: 'my-project'
+      projectId: 'my-project',
+      appId: 'my-appid'
     }
   },
   location: 'us-central1'
@@ -100,4 +101,22 @@ describe('VertexAIModel', () => {
       );
     }
   });
+  it('throws if not passed an app ID', () => {
+    const fakeVertexAI: VertexAI = {
+      app: {
+        name: 'DEFAULT',
+        automaticDataCollectionEnabled: true,
+        options: {
+          apiKey: 'key',
+          projectId: 'my-project'
+        }
+      },
+      location: 'us-central1'
+    };
+    try {
+      new TestModel(fakeVertexAI, 'my-model');
+    } catch (e) {
+      expect((e as VertexAIError).code).to.equal(VertexAIErrorCode.NO_APP_ID);
+    }
+  });
 });

packages/vertexai/src/models/vertexai-model.ts

@@ -68,10 +68,18 @@ export abstract class VertexAIModel {
         VertexAIErrorCode.NO_PROJECT_ID,
         `The "projectId" field is empty in the local Firebase config. Firebase VertexAI requires this field to contain a valid project ID.`
       );
+    } else if (!vertexAI.app?.options?.appId) {
+      throw new VertexAIError(
+        VertexAIErrorCode.NO_APP_ID,
+        `The "appId" field is empty in the local Firebase config. Firebase VertexAI requires this field to contain a valid app ID.`
+      );
     } else {
       this._apiSettings = {
         apiKey: vertexAI.app.options.apiKey,
         project: vertexAI.app.options.projectId,
+        appId: vertexAI.app.options.appId,
+        automaticDataCollectionEnabled:
+          vertexAI.app.automaticDataCollectionEnabled,
         location: vertexAI.location
       };
 

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

@@ -32,6 +32,7 @@ use(chaiAsPromised);
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',
+  appId: 'my-appid',
   location: 'us-central1'
 };
 
@@ -103,6 +104,7 @@ describe('request methods', () => {
     const fakeApiSettings: ApiSettings = {
       apiKey: 'key',
       project: 'myproject',
+      appId: 'my-appid',
       location: 'moon',
       getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
       getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
@@ -124,6 +126,50 @@ describe('request methods', () => {
       const headers = await getHeaders(fakeUrl);
       expect(headers.get('x-goog-api-key')).to.equal('key');
     });
+    it('adds app id if automatedDataCollectionEnabled is true', async () => {
+      const fakeApiSettings: ApiSettings = {
+        apiKey: 'key',
+        project: 'myproject',
+        appId: 'my-appid',
+        location: 'moon',
+        automaticDataCollectionEnabled: true,
+        getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
+        getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
+      };
+      const fakeUrl = new RequestUrl(
+        'models/model-name',
+        Task.GENERATE_CONTENT,
+        fakeApiSettings,
+        true,
+        {}
+      );
+      const headers = await getHeaders(fakeUrl);
+      expect(headers.get('X-Firebase-Appid')).to.equal('my-appid');
+    });
+    it('does not add app id if automatedDataCollectionEnabled is undefined', async () => {
+      const headers = await getHeaders(fakeUrl);
+      expect(headers.get('X-Firebase-Appid')).to.be.null;
+    });
+    it('does not add app id if automatedDataCollectionEnabled is false', async () => {
+      const fakeApiSettings: ApiSettings = {
+        apiKey: 'key',
+        project: 'myproject',
+        appId: 'my-appid',
+        location: 'moon',
+        automaticDataCollectionEnabled: false,
+        getAuthToken: () => Promise.resolve({ accessToken: 'authtoken' }),
+        getAppCheckToken: () => Promise.resolve({ token: 'appchecktoken' })
+      };
+      const fakeUrl = new RequestUrl(
+        'models/model-name',
+        Task.GENERATE_CONTENT,
+        fakeApiSettings,
+        true,
+        {}
+      );
+      const headers = await getHeaders(fakeUrl);
+      expect(headers.get('X-Firebase-Appid')).to.be.null;
+    });
     it('adds app check token if it exists', async () => {
       const headers = await getHeaders(fakeUrl);
       expect(headers.get('X-Firebase-AppCheck')).to.equal('appchecktoken');
@@ -135,6 +181,7 @@ describe('request methods', () => {
         {
           apiKey: 'key',
           project: 'myproject',
+          appId: 'my-appid',
           location: 'moon'
         },
         true,
@@ -167,6 +214,7 @@ describe('request methods', () => {
         {
           apiKey: 'key',
           project: 'myproject',
+          appId: 'my-appid',
           location: 'moon',
           getAppCheckToken: () =>
             Promise.resolve({ token: 'dummytoken', error: Error('oops') })
@@ -193,6 +241,7 @@ describe('request methods', () => {
         {
           apiKey: 'key',
           project: 'myproject',
+          appId: 'my-appid',
           location: 'moon'
         },
         true,

packages/vertexai/src/requests/request.ts

@@ -84,6 +84,9 @@ export async function getHeaders(url: RequestUrl): Promise<Headers> {
   headers.append('Content-Type', 'application/json');
   headers.append('x-goog-api-client', getClientHeaders());
   headers.append('x-goog-api-key', url.apiSettings.apiKey);
+  if (url.apiSettings.automaticDataCollectionEnabled) {
+    headers.append('X-Firebase-Appid', url.apiSettings.appId);
+  }
   if (url.apiSettings.getAppCheckToken) {
     const appCheckToken = await url.apiSettings.getAppCheckToken();
     if (appCheckToken) {

packages/vertexai/src/types/error.ts

@@ -87,6 +87,9 @@ export const enum VertexAIErrorCode {
   /** An error occurred due to a missing Firebase API key. */
   NO_API_KEY = 'no-api-key',
 
+  /** An error occured due to a missing app ID. */
+  NO_APP_ID = 'no-app-id',
+
   /** An error occurred due to a model name not being specified during initialization. */
   NO_MODEL = 'no-model',
 

packages/vertexai/src/types/internal.ts

@@ -23,7 +23,9 @@ export * from './imagen/internal';
 export interface ApiSettings {
   apiKey: string;
   project: string;
+  appId: string;
   location: string;
+  automaticDataCollectionEnabled?: boolean;
   getAuthToken?: () => Promise<FirebaseAuthTokenData | null>;
   getAppCheckToken?: () => Promise<AppCheckTokenResult>;
 }