Add appcheck header (#294)

Author: hsubox76

packages/vertexai/src/api.ts

@@ -19,7 +19,7 @@ import { FirebaseApp, getApp, _getProvider } from '@firebase/app';
 import { Provider } from '@firebase/component';
 import { getModularInstance } from '@firebase/util';
 import { DEFAULT_LOCATION, VERTEX_TYPE } from './constants';
-import { VertexService } from './factory';
+import { VertexService } from './service';
 import { Vertex, VertexOptions } from './public-types';
 import { ERROR_FACTORY, VertexError } from './errors';
 import { ModelParams, RequestOptions } from './types';

packages/vertexai/src/index.ts

@@ -22,7 +22,7 @@
  */
 
 import { registerVersion, _registerComponent } from '@firebase/app';
-import { factory } from './factory';
+import { VertexService } from './service';
 import { VERTEX_TYPE } from './constants';
 import { Component, ComponentType } from '@firebase/component';
 import { name, version } from '../package.json';
@@ -41,7 +41,7 @@ function registerVertex(): void {
         // getImmediate for FirebaseApp will always succeed
         const app = container.getProvider('app').getImmediate();
         const appCheckProvider = container.getProvider('app-check-internal');
-        return factory(app, appCheckProvider, { location });
+        return new VertexService(app, appCheckProvider, { location });
       },
       ComponentType.PUBLIC
     ).setMultipleInstances(true)

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

@@ -20,7 +20,7 @@ import {
   CountTokensResponse,
   RequestOptions
 } from '../types';
-import { RequestUrl, Task, makeRequest } from '../requests/request';
+import { Task, makeRequest } from '../requests/request';
 import { ApiSettings } from '../types/internal';
 
 export async function countTokens(
@@ -29,15 +29,11 @@ export async function countTokens(
   params: CountTokensRequest,
   requestOptions?: RequestOptions
 ): Promise<CountTokensResponse> {
-  const url = new RequestUrl(
+  const response = await makeRequest(
     model,
     Task.COUNT_TOKENS,
     apiSettings,
     false,
-    requestOptions
-  );
-  const response = await makeRequest(
-    url,
     JSON.stringify(params),
     requestOptions
   );

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

@@ -28,6 +28,7 @@ import {
   HarmCategory
 } from '../types';
 import { ApiSettings } from '../types/internal';
+import { Task } from '../requests/request';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -69,10 +70,14 @@ describe('generateContent()', () => {
     );
     expect(result.response.text()).to.include('Helena');
     expect(makeRequestStub).to.be.calledWith(
-      match.instanceOf(request.RequestUrl),
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
       match((value: string) => {
         return value.includes('contents');
-      })
+      }),
+      undefined
     );
   });
   it('long response', async () => {
@@ -88,7 +93,10 @@ describe('generateContent()', () => {
     expect(result.response.text()).to.include('Use Freshly Ground Coffee');
     expect(result.response.text()).to.include('30 minutes of brewing');
     expect(makeRequestStub).to.be.calledWith(
-      match.instanceOf(request.RequestUrl),
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
       match.any
     );
   });
@@ -107,7 +115,10 @@ describe('generateContent()', () => {
       result.response.candidates?.[0].citationMetadata?.citations.length
     ).to.equal(1);
     expect(makeRequestStub).to.be.calledWith(
-      match.instanceOf(request.RequestUrl),
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
       match.any
     );
   });
@@ -125,7 +136,10 @@ describe('generateContent()', () => {
     );
     expect(result.response.text).to.throw('SAFETY');
     expect(makeRequestStub).to.be.calledWith(
-      match.instanceOf(request.RequestUrl),
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
       match.any
     );
   });
@@ -143,7 +157,10 @@ describe('generateContent()', () => {
     );
     expect(result.response.text).to.throw('SAFETY');
     expect(makeRequestStub).to.be.calledWith(
-      match.instanceOf(request.RequestUrl),
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
       match.any
     );
   });
@@ -159,7 +176,10 @@ describe('generateContent()', () => {
     );
     expect(result.response.text()).to.equal('');
     expect(makeRequestStub).to.be.calledWith(
-      match.instanceOf(request.RequestUrl),
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
       match.any
     );
   });
@@ -175,7 +195,10 @@ describe('generateContent()', () => {
     );
     expect(result.response.text()).to.include('30 minutes of brewing');
     expect(makeRequestStub).to.be.calledWith(
-      match.instanceOf(request.RequestUrl),
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
       match.any
     );
   });

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

@@ -22,7 +22,7 @@ import {
   GenerateContentStreamResult,
   RequestOptions
 } from '../types';
-import { RequestUrl, Task, makeRequest } from '../requests/request';
+import { Task, makeRequest } from '../requests/request';
 import { addHelpers } from '../requests/response-helpers';
 import { processStream } from '../requests/stream-reader';
 import { ApiSettings } from '../types/internal';
@@ -33,15 +33,11 @@ export async function generateContentStream(
   params: GenerateContentRequest,
   requestOptions?: RequestOptions
 ): Promise<GenerateContentStreamResult> {
-  const url = new RequestUrl(
+  const response = await makeRequest(
     model,
     Task.STREAM_GENERATE_CONTENT,
     apiSettings,
     /* stream */ true,
-    requestOptions
-  );
-  const response = await makeRequest(
-    url,
     JSON.stringify(params),
     requestOptions
   );
@@ -54,15 +50,11 @@ export async function generateContent(
   params: GenerateContentRequest,
   requestOptions?: RequestOptions
 ): Promise<GenerateContentResult> {
-  const url = new RequestUrl(
+  const response = await makeRequest(
     model,
     Task.GENERATE_CONTENT,
     apiSettings,
     /* stream */ false,
-    requestOptions
-  );
-  const response = await makeRequest(
-    url,
     JSON.stringify(params),
     requestOptions
   );

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

@@ -39,6 +39,7 @@ import { formatGenerateContentInput } from '../requests/request-helpers';
 import { Vertex } from '../public-types';
 import { ERROR_FACTORY, VertexError } from '../errors';
 import { ApiSettings } from '../types/internal';
+import { VertexService } from '../service';
 
 /**
  * Class for generative model APIs.
@@ -67,6 +68,10 @@ export class GenerativeModel {
         project: vertex.app.options.projectId,
         location: vertex.location
       };
+      if ((vertex as VertexService).appCheck) {
+        this._apiSettings.getAppCheckToken = () =>
+          (vertex as VertexService).appCheck!.getToken();
+      }
     }
     if (modelParams.model.includes('/')) {
       if (modelParams.model.startsWith('models/')) {