diff --git a/.changeset/nasty-rings-drop.md b/.changeset/nasty-rings-drop.md new file mode 100644 index 00000000000..7755dced9f1 --- /dev/null +++ b/.changeset/nasty-rings-drop.md @@ -0,0 +1,6 @@ +--- +'@firebase/ai': minor +'firebase': minor +--- + +Add App Check limited use token option to `getAI()`. diff --git a/common/api-review/ai.api.md b/common/api-review/ai.api.md index e0eac35996a..5f10f453308 100644 --- a/common/api-review/ai.api.md +++ b/common/api-review/ai.api.md @@ -4,10 +4,18 @@ ```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 { @@ -15,6 +23,7 @@ export interface AI { backend: Backend; // @deprecated (undocumented) location: string; + options?: AIOptions; } // @public @@ -53,7 +62,7 @@ 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; @@ -61,7 +70,8 @@ export abstract class AIModel { // @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) diff --git a/docs-devsite/ai.ai.md b/docs-devsite/ai.ai.md index d4127ffb7e8..d7aafb6a5b5 100644 --- a/docs-devsite/ai.ai.md +++ b/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. + +Signature: + +```typescript +options?: AIOptions; +``` diff --git a/docs-devsite/ai.aioptions.md b/docs-devsite/ai.aioptions.md index a092046900b..a5b326ef004 100644 --- a/docs-devsite/ai.aioptions.md +++ b/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)). Signature: ```typescript -backend: Backend; +backend?: Backend; +``` + +## AIOptions.useLimitedUseAppCheckTokens + +Whether to use App Check limited use tokens. Defaults to false. + +Signature: + +```typescript +useLimitedUseAppCheckTokens?: boolean; ``` diff --git a/docs-devsite/ai.md b/docs-devsite/ai.md index 29b3f73f86e..3c7669204a1 100644 --- a/docs-devsite/ai.md +++ b/docs-devsite/ai.md @@ -21,6 +21,8 @@ The Firebase AI Web SDK. | function(ai, ...) | | [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) | (Public Preview) Returns an [ImagenModel](./ai.imagenmodel.md#imagenmodel_class) class with methods for using Imagen.Only Imagen 3 models (named imagen-3.0-*) are supported. | +| function(container, ...) | +| [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} + +Signature: + +```typescript +export declare function factory(container: ComponentContainer, { instanceIdentifier }: InstanceFactoryOptions): AIService; +``` + +#### Parameters + +| Parameter | Type | Description | +| --- | --- | --- | +| container | ComponentContainer | | +| { instanceIdentifier } | InstanceFactoryOptions | | + +Returns: + +AIService + ## AIErrorCode Standardized error codes that [AIError](./ai.aierror.md#aierror_class) can have. diff --git a/packages/ai/src/api.test.ts b/packages/ai/src/api.test.ts index 76a9b4523c2..55d2eaa4ad3 100644 --- a/packages/ai/src/api.test.ts +++ b/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); diff --git a/packages/ai/src/api.ts b/packages/ai/src/api.ts index 62c7c27f07a..26d2904150d 100644 --- a/packages/ai/src/api.ts +++ b/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 = { + useLimitedUseAppCheckTokens: options?.useLimitedUseAppCheckTokens ?? false + }; + + const identifier = encodeInstanceIdentifier(backend); + const aiInstance = AIProvider.getImmediate({ identifier }); + + aiInstance.options = finalOptions; + + return aiInstance; } /** diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 6ad1f2e3f08..8ac9db61b6d 100644 --- a/packages/ai/src/index.ts +++ b/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); diff --git a/packages/ai/src/models/ai-model.test.ts b/packages/ai/src/models/ai-model.test.ts index 4f23fe9d06f..2e8f8998c58 100644 --- a/packages/ai/src/models/ai-model.test.ts +++ b/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: { diff --git a/packages/ai/src/models/ai-model.ts b/packages/ai/src/models/ai-model.ts index 084dbe329cc..3fe202d5eb2 100644 --- a/packages/ai/src/models/ai-model.ts +++ b/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) { diff --git a/packages/ai/src/public-types.ts b/packages/ai/src/public-types.ts index 57812f20c1e..fff41251a01 100644 --- a/packages/ai/src/public-types.ts +++ b/packages/ai/src/public-types.ts @@ -38,6 +38,10 @@ export interface AI { * Vertex AI Gemini API (using {@link VertexAIBackend}). */ backend: Backend; + /** + * Options applied to this {@link AI} instance. + */ + options?: AIOptions; /** * @deprecated use `AI.backend.location` instead. * @@ -90,6 +94,11 @@ export type BackendType = (typeof BackendType)[keyof typeof BackendType]; export interface AIOptions { /** * The backend configuration to use for the AI service instance. + * Defaults to the Gemini Developer API backend ({@link GoogleAIBackend}). */ - backend: Backend; + backend?: Backend; + /** + * Whether to use App Check limited use tokens. Defaults to false. + */ + useLimitedUseAppCheckTokens?: boolean; } diff --git a/packages/ai/src/service.ts b/packages/ai/src/service.ts index 006cc45a94e..c5c9d9a036d 100644 --- a/packages/ai/src/service.ts +++ b/packages/ai/src/service.ts @@ -16,7 +16,7 @@ */ import { FirebaseApp, _FirebaseService } from '@firebase/app'; -import { AI } from './public-types'; +import { AI, AIOptions } from './public-types'; import { AppCheckInternalComponentName, FirebaseAppCheckInternal @@ -31,6 +31,7 @@ import { Backend, VertexAIBackend } from './backend'; export class AIService implements AI, _FirebaseService { auth: FirebaseAuthInternal | null; appCheck: FirebaseAppCheckInternal | null; + _options?: Omit; location: string; // This is here for backwards-compatibility constructor( @@ -54,4 +55,12 @@ export class AIService implements AI, _FirebaseService { _delete(): Promise { return Promise.resolve(); } + + set options(optionsToSet: AIOptions) { + this._options = optionsToSet; + } + + get options(): AIOptions | undefined { + return this._options; + } } diff --git a/packages/ai/test-utils/get-fake-firebase-services.ts b/packages/ai/test-utils/get-fake-firebase-services.ts new file mode 100644 index 00000000000..20ae7fb70be --- /dev/null +++ b/packages/ai/test-utils/get-fake-firebase-services.ts @@ -0,0 +1,67 @@ +/** + * @license + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + FirebaseApp, + initializeApp, + _registerComponent, + _addOrOverwriteComponent +} from '@firebase/app'; +import { Component, ComponentType } from '@firebase/component'; +import { FirebaseAppCheckInternal } from '@firebase/app-check-interop-types'; +import { AI_TYPE } from '../src/constants'; +import { factory } from '../src'; + +const fakeConfig = { + projectId: 'projectId', + authDomain: 'authDomain', + messagingSenderId: 'messagingSenderId', + databaseURL: 'databaseUrl', + storageBucket: 'storageBucket' +}; + +export function getFullApp(fakeAppParams?: { + appId?: string; + apiKey?: string; +}): FirebaseApp { + _registerComponent(new Component(AI_TYPE, factory, ComponentType.PUBLIC)); + _registerComponent( + new Component( + 'app-check-internal', + () => { + return {} as FirebaseAppCheckInternal; + }, + ComponentType.PUBLIC + ) + ); + const app = initializeApp({ ...fakeConfig, ...fakeAppParams }); + _addOrOverwriteComponent( + app, + //@ts-ignore + new Component( + 'heartbeat', + // @ts-ignore + () => { + return { + triggerHeartbeat: () => {} + }; + }, + ComponentType.PUBLIC + ) + ); + return app; +}