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;
+    appCheck?: AppCheckOptions;
+    backend?: Backend;
 }
 
 // @public
@@ -75,6 +85,11 @@ export class AnyOfSchema extends Schema {
     toJSON(): SchemaRequest;
 }
 
+// @public
+export interface AppCheckOptions {
+    limitedUseTokens?: boolean;
+}
+
 // @public
 export class ArraySchema extends Schema {
     constructor(schemaParams: SchemaParams, items: TypedSchema);
@@ -229,6 +244,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/_toc.yaml

@@ -14,6 +14,8 @@ toc:
     path: /docs/reference/js/ai.aioptions.md
   - title: AnyOfSchema
     path: /docs/reference/js/ai.anyofschema.md
+  - title: AppCheckOptions
+    path: /docs/reference/js/ai.appcheckoptions.md
   - title: ArraySchema
     path: /docs/reference/js/ai.arrayschema.md
   - title: Backend

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. |
+|  [appCheck](./ai.aioptions.md#aioptionsappcheck) | [AppCheckOptions](./ai.appcheckoptions.md#appcheckoptions_interface) | Configures App Check usage for this 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 [GoogleAIBackend](./ai.googleaibackend.md#googleaibackend_class)<!-- -->. |
+
+## AIOptions.appCheck
+
+Configures App Check usage for this AI service instance.
+
+<b>Signature:</b>
+
+```typescript
+appCheck?: AppCheckOptions;
+```
 
 ## AIOptions.backend
 
-The backend configuration to use for the AI service instance.
+The backend configuration to use for the AI service instance. Defaults to [GoogleAIBackend](./ai.googleaibackend.md#googleaibackend_class)<!-- -->.
 
 <b>Signature:</b>
 
 ```typescript
-backend: Backend;
+backend?: Backend;
 ```

docs-devsite/ai.appcheckoptions.md

@@ -0,0 +1,35 @@
+Project: /docs/reference/js/_project.yaml
+Book: /docs/reference/_book.yaml
+page_type: reference
+
+{% comment %}
+DO NOT EDIT THIS FILE!
+This is generated by the JS SDK team, and any local changes will be
+overwritten. Changes should be made in the source code at
+https://github.com/firebase/firebase-js-sdk
+{% endcomment %}
+
+# AppCheckOptions interface
+Configures App Check usage for this AI service instance.
+
+<b>Signature:</b>
+
+```typescript
+export interface AppCheckOptions 
+```
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [limitedUseTokens](./ai.appcheckoptions.md#appcheckoptionslimitedusetokens) | boolean | Defaults to false. |
+
+## AppCheckOptions.limitedUseTokens
+
+Defaults to false.
+
+<b>Signature:</b>
+
+```typescript
+limitedUseTokens?: 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
 
@@ -50,6 +52,7 @@ The Firebase AI Web SDK.
 |  --- | --- |
 |  [AI](./ai.ai.md#ai_interface) | An instance of the Firebase AI SDK.<!-- -->Do not create this instance directly. Instead, use [getAI()](./ai.md#getai_a94a413)<!-- -->. |
 |  [AIOptions](./ai.aioptions.md#aioptions_interface) | Options for initializing the AI service using [getAI()](./ai.md#getai_a94a413)<!-- -->. This allows specifying which backend to use (Vertex AI Gemini API or Gemini Developer API) and configuring its specific options (like location for Vertex AI). |
+|  [AppCheckOptions](./ai.appcheckoptions.md#appcheckoptions_interface) | Configures App Check usage for this AI service instance. |
 |  [BaseParams](./ai.baseparams.md#baseparams_interface) | Base parameters for a number of methods. |
 |  [ChromeAdapter](./ai.chromeadapter.md#chromeadapter_interface) | <b>(EXPERIMENTAL)</b> Defines an inference "backend" that uses Chrome's on-device model, and encapsulates logic for detecting when on-device inference is possible.<!-- -->These methods should not be called directly by the user. |
 |  [Citation](./ai.citation.md#citation_interface) | A single citation. |
@@ -278,6 +281,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(), { appCheck: { limitedUseTokens: true } });
+      expect(ai.backend).to.be.instanceOf(GoogleAIBackend);
+      expect(ai.options?.appCheck?.limitedUseTokens).to.be.true;
+    });
+    it('works with options: backend specified, limited use token', () => {
+      const ai = getAI(getFullApp(), {
+        backend: new VertexAIBackend('us-central1'),
+        appCheck: { limitedUseTokens: true }
+      });
+      expect(ai.backend).to.be.instanceOf(VertexAIBackend);
+      expect(ai.options?.appCheck?.limitedUseTokens).to.be.true;
+    });
+    it('works with options: appCheck option is empty', () => {
+      const ai = getAI(getFullApp(), {
+        backend: new VertexAIBackend('us-central1'),
+        appCheck: {}
+      });
+      expect(ai.backend).to.be.instanceOf(VertexAIBackend);
+      expect(ai.options?.appCheck?.limitedUseTokens).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?.appCheck?.limitedUseTokens).to.be.false;
+    });
+  });
   it('getGenerativeModel throws if no model is provided', () => {
     try {
       getGenerativeModel(fakeAI, {} as ModelParams);

packages/ai/src/api.ts

@@ -75,18 +75,33 @@ 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'> = {};
+
+  if (options?.appCheck) {
+    if (options.appCheck.hasOwnProperty('limitedUseTokens')) {
+      finalOptions.appCheck = options.appCheck;
+    } else {
+      finalOptions.appCheck = { ...options.appCheck, limitedUseTokens: false };
+    }
+  } else {
+    finalOptions.appCheck = { limitedUseTokens: 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);