Add limitedUseToken option to AI SDK

.changeset/nasty-rings-drop.md

@@ -0,0 +1,6 @@
+---
+'@firebase/ai': minor
+'firebase': minor
+---
+
+Add App Check limited use token option to `getAI()`.

common/api-review/ai.api.md

@@ -4,17 +4,26 @@
 
 ```ts
 
+import { AppCheckInternalComponentName } from '@firebase/app-check-interop-types';
 import { AppCheckTokenResult } from '@firebase/app-check-interop-types';
+import { ComponentContainer } from '@firebase/component';
 import { FirebaseApp } from '@firebase/app';
+import { FirebaseAppCheckInternal } from '@firebase/app-check-interop-types';
+import { FirebaseAuthInternal } from '@firebase/auth-interop-types';
+import { FirebaseAuthInternalName } from '@firebase/auth-interop-types';
 import { FirebaseAuthTokenData } from '@firebase/auth-interop-types';
 import { FirebaseError } from '@firebase/util';
+import { _FirebaseService } from '@firebase/app';
+import { InstanceFactoryOptions } from '@firebase/component';
+import { Provider } from '@firebase/component';
 
 // @public
 export interface AI {
     app: FirebaseApp;
     backend: Backend;
     // @deprecated (undocumented)
     location: string;
+    options?: AIOptions;
 }
 
 // @public
@@ -53,15 +62,16 @@ export abstract class AIModel {
     // Warning: (ae-forgotten-export) The symbol "ApiSettings" needs to be exported by the entry point index.d.ts
     //
     // @internal (undocumented)
-    protected _apiSettings: ApiSettings;
+    _apiSettings: ApiSettings;
     readonly model: string;
     // @internal
     static normalizeModelName(modelName: string, backendType: BackendType): string;
     }
 
 // @public
 export interface AIOptions {
-    backend: Backend;
+    backend?: Backend;
+    useLimitedUseAppCheckTokens?: boolean;
 }
 
 // @public
@@ -229,6 +239,11 @@ export interface ErrorDetails {
     reason?: string;
 }
 
+// Warning: (ae-forgotten-export) The symbol "AIService" needs to be exported by the entry point index.d.ts
+//
+// @public (undocumented)
+export function factory(container: ComponentContainer, { instanceIdentifier }: InstanceFactoryOptions): AIService;
+
 // @public
 export interface FileData {
     // (undocumented)

docs-devsite/ai.ai.md

@@ -27,6 +27,7 @@ export interface AI
 |  [app](./ai.ai.md#aiapp) | [FirebaseApp](./app.firebaseapp.md#firebaseapp_interface) | The [FirebaseApp](./app.firebaseapp.md#firebaseapp_interface) this [AI](./ai.ai.md#ai_interface) instance is associated with. |
 |  [backend](./ai.ai.md#aibackend) | [Backend](./ai.backend.md#backend_class) | A [Backend](./ai.backend.md#backend_class) instance that specifies the configuration for the target backend, either the Gemini Developer API (using [GoogleAIBackend](./ai.googleaibackend.md#googleaibackend_class)<!-- -->) or the Vertex AI Gemini API (using [VertexAIBackend](./ai.vertexaibackend.md#vertexaibackend_class)<!-- -->). |
 |  [location](./ai.ai.md#ailocation) | string |  |
+|  [options](./ai.ai.md#aioptions) | [AIOptions](./ai.aioptions.md#aioptions_interface) | Options applied to this [AI](./ai.ai.md#ai_interface) instance. |
 
 ## AI.app
 
@@ -62,3 +63,13 @@ backend: Backend;
 ```typescript
 location: string;
 ```
+
+## AI.options
+
+Options applied to this [AI](./ai.ai.md#ai_interface) instance.
+
+<b>Signature:</b>
+
+```typescript
+options?: AIOptions;
+```

docs-devsite/ai.aioptions.md

@@ -22,14 +22,25 @@ export interface AIOptions
 
 |  Property | Type | Description |
 |  --- | --- | --- |
-|  [backend](./ai.aioptions.md#aioptionsbackend) | [Backend](./ai.backend.md#backend_class) | The backend configuration to use for the AI service instance. |
+|  [backend](./ai.aioptions.md#aioptionsbackend) | [Backend](./ai.backend.md#backend_class) | The backend configuration to use for the AI service instance. Defaults to the Gemini Developer API backend ([GoogleAIBackend](./ai.googleaibackend.md#googleaibackend_class)<!-- -->). |
+|  [useLimitedUseAppCheckTokens](./ai.aioptions.md#aioptionsuselimiteduseappchecktokens) | boolean | Whether to use App Check limited use tokens. Defaults to false. |
 
 ## AIOptions.backend
 
-The backend configuration to use for the AI service instance.
+The backend configuration to use for the AI service instance. Defaults to the Gemini Developer API backend ([GoogleAIBackend](./ai.googleaibackend.md#googleaibackend_class)<!-- -->).
 
 <b>Signature:</b>
 
 ```typescript
-backend: Backend;
+backend?: Backend;
+```
+
+## AIOptions.useLimitedUseAppCheckTokens
+
+Whether to use App Check limited use tokens. Defaults to false.
+
+<b>Signature:</b>
+
+```typescript
+useLimitedUseAppCheckTokens?: boolean;
 ```

docs-devsite/ai.md

@@ -21,6 +21,8 @@ The Firebase AI Web SDK.
 |  <b>function(ai, ...)</b> |
 |  [getGenerativeModel(ai, modelParams, requestOptions)](./ai.md#getgenerativemodel_c63f46a) | Returns a [GenerativeModel](./ai.generativemodel.md#generativemodel_class) class with methods for inference and other functionality. |
 |  [getImagenModel(ai, modelParams, requestOptions)](./ai.md#getimagenmodel_e1f6645) | <b><i>(Public Preview)</i></b> Returns an [ImagenModel](./ai.imagenmodel.md#imagenmodel_class) class with methods for using Imagen.<!-- -->Only Imagen 3 models (named <code>imagen-3.0-*</code>) are supported. |
+|  <b>function(container, ...)</b> |
+|  [factory(container, { instanceIdentifier })](./ai.md#factory_6581aeb) |  |
 
 ## Classes
 
@@ -278,6 +280,27 @@ export declare function getImagenModel(ai: AI, modelParams: ImagenModelParams, r
 
 If the `apiKey` or `projectId` fields are missing in your Firebase config.
 
+## function(container, ...)
+
+### factory(container, { instanceIdentifier }) {:#factory_6581aeb}
+
+<b>Signature:</b>
+
+```typescript
+export declare function factory(container: ComponentContainer, { instanceIdentifier }: InstanceFactoryOptions): AIService;
+```
+
+#### Parameters
+
+|  Parameter | Type | Description |
+|  --- | --- | --- |
+|  container | ComponentContainer |  |
+|  { instanceIdentifier } | InstanceFactoryOptions |  |
+
+<b>Returns:</b>
+
+AIService
+
 ## AIErrorCode
 
 Standardized error codes that [AIError](./ai.aierror.md#aierror_class) can have.

packages/ai/src/api.test.ts

@@ -16,11 +16,12 @@
  */
 import { ImagenModelParams, ModelParams, AIErrorCode } from './types';
 import { AIError } from './errors';
-import { ImagenModel, getGenerativeModel, getImagenModel } from './api';
+import { getAI, ImagenModel, getGenerativeModel, getImagenModel } from './api';
 import { expect } from 'chai';
 import { AI } from './public-types';
 import { GenerativeModel } from './models/generative-model';
-import { VertexAIBackend } from './backend';
+import { GoogleAIBackend, VertexAIBackend } from './backend';
+import { getFullApp } from '../test-utils/get-fake-firebase-services';
 import { AI_TYPE, DEFAULT_HYBRID_IN_CLOUD_MODEL } from './constants';
 
 const fakeAI: AI = {
@@ -38,6 +39,40 @@ const fakeAI: AI = {
 };
 
 describe('Top level API', () => {
+  describe('getAI()', () => {
+    it('works without options', () => {
+      const ai = getAI(getFullApp());
+      expect(ai.backend).to.be.instanceOf(GoogleAIBackend);
+    });
+    it('works with options: no backend, limited use token', () => {
+      const ai = getAI(getFullApp(), { useLimitedUseAppCheckTokens: true });
+      expect(ai.backend).to.be.instanceOf(GoogleAIBackend);
+      expect(ai.options?.useLimitedUseAppCheckTokens).to.be.true;
+    });
+    it('works with options: backend specified, limited use token', () => {
+      const ai = getAI(getFullApp(), {
+        backend: new VertexAIBackend('us-central1'),
+        useLimitedUseAppCheckTokens: true
+      });
+      expect(ai.backend).to.be.instanceOf(VertexAIBackend);
+      expect(ai.options?.useLimitedUseAppCheckTokens).to.be.true;
+    });
+    it('works with options: appCheck option is falsy', () => {
+      const ai = getAI(getFullApp(), {
+        backend: new VertexAIBackend('us-central1'),
+        useLimitedUseAppCheckTokens: undefined
+      });
+      expect(ai.backend).to.be.instanceOf(VertexAIBackend);
+      expect(ai.options?.useLimitedUseAppCheckTokens).to.be.false;
+    });
+    it('works with options: backend specified only', () => {
+      const ai = getAI(getFullApp(), {
+        backend: new VertexAIBackend('us-central1')
+      });
+      expect(ai.backend).to.be.instanceOf(VertexAIBackend);
+      expect(ai.options?.useLimitedUseAppCheckTokens).to.be.false;
+    });
+  });
   it('getGenerativeModel throws if no model is provided', () => {
     try {
       getGenerativeModel(fakeAI, {} as ModelParams);

packages/ai/src/api.ts

@@ -75,18 +75,25 @@ declare module '@firebase/component' {
  *
  * @public
  */
-export function getAI(
-  app: FirebaseApp = getApp(),
-  options: AIOptions = { backend: new GoogleAIBackend() }
-): AI {
+export function getAI(app: FirebaseApp = getApp(), options?: AIOptions): AI {
   app = getModularInstance(app);
   // Dependencies
   const AIProvider: Provider<'AI'> = _getProvider(app, AI_TYPE);
 
-  const identifier = encodeInstanceIdentifier(options.backend);
-  return AIProvider.getImmediate({
+  const backend = options?.backend ?? new GoogleAIBackend();
+
+  const finalOptions: Omit<AIOptions, 'backend'> = {
+    useLimitedUseAppCheckTokens: options?.useLimitedUseAppCheckTokens ?? false
+  };
+
+  const identifier = encodeInstanceIdentifier(backend);
+  const aiInstance = AIProvider.getImmediate({
     identifier
   });
+
+  aiInstance.options = finalOptions;
+
+  return aiInstance;
 }
 
 /**

packages/ai/src/index.ts

@@ -24,7 +24,12 @@
 import { registerVersion, _registerComponent } from '@firebase/app';
 import { AIService } from './service';
 import { AI_TYPE } from './constants';
-import { Component, ComponentType } from '@firebase/component';
+import {
+  Component,
+  ComponentContainer,
+  ComponentType,
+  InstanceFactoryOptions
+} from '@firebase/component';
 import { name, version } from '../package.json';
 import { decodeInstanceIdentifier } from './helpers';
 import { AIError } from './api';
@@ -36,28 +41,31 @@ declare global {
   }
 }
 
-function registerAI(): void {
-  _registerComponent(
-    new Component(
-      AI_TYPE,
-      (container, { instanceIdentifier }) => {
-        if (!instanceIdentifier) {
-          throw new AIError(
-            AIErrorCode.ERROR,
-            'AIService instance identifier is undefined.'
-          );
-        }
+export function factory(
+  container: ComponentContainer,
+  { instanceIdentifier }: InstanceFactoryOptions
+): AIService {
+  if (!instanceIdentifier) {
+    throw new AIError(
+      AIErrorCode.ERROR,
+      'AIService instance identifier is undefined.'
+    );
+  }
 
-        const backend = decodeInstanceIdentifier(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
-    ).setMultipleInstances(true)
+  // 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);
+}
+
+function registerAI(): void {
+  _registerComponent(
+    new Component(AI_TYPE, factory, ComponentType.PUBLIC).setMultipleInstances(
+      true
+    )
   );
 
   registerVersion(name, version);

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

@@ -17,9 +17,11 @@
 import { use, expect } from 'chai';
 import { AI, AIErrorCode } from '../public-types';
 import sinonChai from 'sinon-chai';
+import { stub } from 'sinon';
 import { AIModel } from './ai-model';
 import { AIError } from '../errors';
 import { VertexAIBackend } from '../backend';
+import { AIService } from '../service';
 
 use(sinonChai);
 
@@ -67,6 +69,52 @@ describe('AIModel', () => {
     const testModel = new TestModel(fakeAI, 'tunedModels/my-model');
     expect(testModel.model).to.equal('tunedModels/my-model');
   });
+  it('calls regular app check token when option is set', async () => {
+    const getTokenStub = stub().resolves();
+    const getLimitedUseTokenStub = stub().resolves();
+    const testModel = new TestModel(
+      //@ts-ignore
+      {
+        ...fakeAI,
+        options: { useLimitedUseAppCheckTokens: false },
+        appCheck: {
+          getToken: getTokenStub,
+          getLimitedUseToken: getLimitedUseTokenStub
+        }
+      } as AIService,
+      'models/my-model'
+    );
+    if (testModel._apiSettings?.getAppCheckToken) {
+      await testModel._apiSettings.getAppCheckToken();
+    }
+    expect(getTokenStub).to.be.called;
+    expect(getLimitedUseTokenStub).to.not.be.called;
+    getTokenStub.reset();
+    getLimitedUseTokenStub.reset();
+  });
+  it('calls limited use token when option is set', async () => {
+    const getTokenStub = stub().resolves();
+    const getLimitedUseTokenStub = stub().resolves();
+    const testModel = new TestModel(
+      //@ts-ignore
+      {
+        ...fakeAI,
+        options: { useLimitedUseAppCheckTokens: true },
+        appCheck: {
+          getToken: getTokenStub,
+          getLimitedUseToken: getLimitedUseTokenStub
+        }
+      } as AIService,
+      'models/my-model'
+    );
+    if (testModel._apiSettings?.getAppCheckToken) {
+      await testModel._apiSettings.getAppCheckToken();
+    }
+    expect(getTokenStub).to.not.be.called;
+    expect(getLimitedUseTokenStub).to.be.called;
+    getTokenStub.reset();
+    getLimitedUseTokenStub.reset();
+  });
   it('throws if not passed an api key', () => {
     const fakeAI: AI = {
       app: {

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

@@ -39,7 +39,7 @@ export abstract class AIModel {
   /**
    * @internal
    */
-  protected _apiSettings: ApiSettings;
+  _apiSettings: ApiSettings;
 
   /**
    * Constructs a new instance of the {@link AIModel} class.
@@ -90,8 +90,13 @@ export abstract class AIModel {
           return Promise.resolve({ token });
         };
       } else if ((ai as AIService).appCheck) {
-        this._apiSettings.getAppCheckToken = () =>
-          (ai as AIService).appCheck!.getToken();
+        if (ai.options?.useLimitedUseAppCheckTokens) {
+          this._apiSettings.getAppCheckToken = () =>
+            (ai as AIService).appCheck!.getLimitedUseToken();
+        } else {
+          this._apiSettings.getAppCheckToken = () =>
+            (ai as AIService).appCheck!.getToken();
+        }
       }
 
       if ((ai as AIService).auth) {