Only send if data collection is on

Author: dlarocque

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

@@ -78,6 +78,8 @@ export abstract class VertexAIModel {
         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

@@ -126,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 undefined', async () => {
+      const headers = await getHeaders(fakeUrl);
+      expect(headers.get('X-Firebase-AppId')).to.equal('my-appid');
+    });
+    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 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');

packages/vertexai/src/requests/request.ts

@@ -84,7 +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);
-  headers.append('X-Firebase-AppId', url.apiSettings.appId); // Will be converted to 'X-Firebase-Appid' before it's sent in the browser.
+  if (url.apiSettings.automaticDataCollectionEnabled !== false) {
+    headers.append('X-Firebase-AppId', url.apiSettings.appId);
+  }
   if (url.apiSettings.getAppCheckToken) {
     const appCheckToken = await url.apiSettings.getAppCheckToken();
     if (appCheckToken) {

packages/vertexai/src/types/internal.ts

@@ -25,6 +25,7 @@ export interface ApiSettings {
   project: string;
   appId: string;
   location: string;
+  automaticDataCollectionEnabled?: boolean;
   getAuthToken?: () => Promise<FirebaseAuthTokenData | null>;
   getAppCheckToken?: () => Promise<AppCheckTokenResult>;
 }