feat(ai)!: Add explicit initializeDeviceModel() method when using hybrid (#9800)

2026-09-24 16:59:10 +00:00 · 2026-05-26 10:04:04 -07:00
parent 7b87b29ef7
commit 013797f6df
15 changed files with 450 additions and 72 deletions
@@ -0,0 +1,6 @@
+---
+'@firebase/ai': minor
+'firebase': minor
+---
+
+Add `initializeDeviceModel()` method to allow explicit initialization of on-device model in hybrid mode.
@@ -603,6 +603,7 @@ export class GenerativeModel extends AIModel {
    generateContentStream(request: GenerateContentRequest | string | Array<string | Part>, singleRequestOptions?: SingleRequestOptions): Promise<GenerateContentStreamResult>;
    // (undocumented)
    generationConfig: GenerationConfig;
+    initializeDeviceModel(onDownloadProgress?: (progressValue: number) => void): Promise<void>;
    // (undocumented)
    requestOptions?: RequestOptions;
    // (undocumented)
@@ -1002,8 +1003,9 @@ export type Language = (typeof Language)[keyof typeof Language];

 // @beta
 export interface LanguageModelCreateCoreOptions {
-    // (undocumented)
    expectedInputs?: LanguageModelExpected[];
+    expectedOutputs?: LanguageModelExpected[];
+    monitor?: (monitor: LanguageModelDownloadMonitor) => void;
    // @deprecated (undocumented)
    temperature?: number;
    // @deprecated (undocumented)
@@ -1018,6 +1020,18 @@ export interface LanguageModelCreateOptions extends LanguageModelCreateCoreOptio
    signal?: AbortSignal;
 }

+// @beta
+export interface LanguageModelDownloadMonitor {
+    // (undocumented)
+    addEventListener: (eventType: 'downloadprogress', eventListener: (e: {
+        loaded: number;
+    }) => void) => void;
+    // (undocumented)
+    removeEventListener: (eventType: 'downloadprogress', eventListener: (e: {
+        loaded: number;
+    }) => void) => void;
+}
+
 // @beta
 export interface LanguageModelExpected {
    // (undocumented)
@@ -142,6 +142,8 @@ toc:
    path: /docs/reference/js/ai.languagemodelcreatecoreoptions.md
  - title: LanguageModelCreateOptions
    path: /docs/reference/js/ai.languagemodelcreateoptions.md
+  - title: LanguageModelDownloadMonitor
+    path: /docs/reference/js/ai.languagemodeldownloadmonitor.md
  - title: LanguageModelExpected
    path: /docs/reference/js/ai.languagemodelexpected.md
  - title: LanguageModelMessage
@@ -43,6 +43,7 @@ export declare class GenerativeModel extends AIModel
 |  [countTokens(request, singleRequestOptions)](./ai.generativemodel.md#generativemodelcounttokens) |  | Counts the tokens in the provided request. |
 |  [generateContent(request, singleRequestOptions)](./ai.generativemodel.md#generativemodelgeneratecontent) |  | Makes a single non-streaming call to the model and returns an object containing a single [GenerateContentResponse](./ai.generatecontentresponse.md#generatecontentresponse_interface)<!-- -->. |
 |  [generateContentStream(request, singleRequestOptions)](./ai.generativemodel.md#generativemodelgeneratecontentstream) |  | Makes a single streaming call to the model and returns an object containing an iterable stream that iterates over all chunks in the streaming response as well as a promise that returns the final aggregated response. |
+|  [initializeDeviceModel(onDownloadProgress)](./ai.generativemodel.md#generativemodelinitializedevicemodel) |  | Initializes on-device models. |
 |  [startChat(startChatParams)](./ai.generativemodel.md#generativemodelstartchat) |  | Gets a new [ChatSession](./ai.chatsession.md#chatsession_class) instance which can be used for multi-turn chats. |

 ## GenerativeModel.(constructor)
@@ -175,6 +176,32 @@ generateContentStream(request: GenerateContentRequest | string | Array<string |

 Promise&lt;[GenerateContentStreamResult](./ai.generatecontentstreamresult.md#generatecontentstreamresult_interface)<!-- -->&gt;

+## GenerativeModel.initializeDeviceModel()
+
+Initializes on-device models.
+
+This may trigger a download on first use. Wait for this promise to complete before calling inference methods if you want to ensure the device models are ready before any calls. Calling inference methods before the device is ready will result in a cloud fallback if `inferenceMode` is set to `PREFER_ON_DEVICE`<!-- -->, and an error if set to `ONLY_ON_DEVICE`<!-- -->.
+
+IMPORTANT: This call must be made on or after a user has interacted with the page (for example, through a button click or key press). If it is called without a user interaction, and it requires a download, this will cause an error.
+
+See the [Prompt API docs](https://developer.chrome.com/docs/ai/prompt-api#use_the_prompt_api) for more details on this requirement.
+
+<b>Signature:</b>
+
+```typescript
+initializeDeviceModel(onDownloadProgress?: (progressValue: number) => void): Promise<void>;
+```
+
+#### Parameters
+
+|  Parameter | Type | Description |
+|  --- | --- | --- |
+|  onDownloadProgress | (progressValue: number) =&gt; void | A callback called repeatedly as the download progresses that provides a <code>progressValue</code> between 0 and 1 representing how much of the download is complete. This will be ignored if <code>monitor</code> was populated in [LanguageModelCreateOptions](./ai.languagemodelcreateoptions.md#languagemodelcreateoptions_interface)<!-- -->. |
+
+<b>Returns:</b>
+
+Promise&lt;void&gt;
+
 ## GenerativeModel.startChat()

 Gets a new [ChatSession](./ai.chatsession.md#chatsession_class) instance which can be used for multi-turn chats.
@@ -25,7 +25,9 @@ export interface LanguageModelCreateCoreOptions

 |  Property | Type | Description |
 |  --- | --- | --- |
-|  [expectedInputs](./ai.languagemodelcreatecoreoptions.md#languagemodelcreatecoreoptionsexpectedinputs) | [LanguageModelExpected](./ai.languagemodelexpected.md#languagemodelexpected_interface)<!-- -->\[\] | <b><i>(Public Preview)</i></b> |
+|  [expectedInputs](./ai.languagemodelcreatecoreoptions.md#languagemodelcreatecoreoptionsexpectedinputs) | [LanguageModelExpected](./ai.languagemodelexpected.md#languagemodelexpected_interface)<!-- -->\[\] | <b><i>(Public Preview)</i></b> Defaults to image input and English text input if not overriden. |
+|  [expectedOutputs](./ai.languagemodelcreatecoreoptions.md#languagemodelcreatecoreoptionsexpectedoutputs) | [LanguageModelExpected](./ai.languagemodelexpected.md#languagemodelexpected_interface)<!-- -->\[\] | <b><i>(Public Preview)</i></b> Defaults to English text output if not overriden. |
+|  [monitor](./ai.languagemodelcreatecoreoptions.md#languagemodelcreatecoreoptionsmonitor) | (monitor: [LanguageModelDownloadMonitor](./ai.languagemodeldownloadmonitor.md#languagemodeldownloadmonitor_interface)<!-- -->) =&gt; void | <b><i>(Public Preview)</i></b> The native download monitor provided by the Chrome API. This provides direct access to the monitor object and is only needed for advanced cases. For a simpler download monitoring interface, provide an <code>onDownloadProgress</code> callback to [GenerativeModel.initializeDeviceModel()](./ai.generativemodel.md#generativemodelinitializedevicemodel) instead. |
 |  [temperature](./ai.languagemodelcreatecoreoptions.md#languagemodelcreatecoreoptionstemperature) | number | <b><i>(Public Preview)</i></b> |
 |  [topK](./ai.languagemodelcreatecoreoptions.md#languagemodelcreatecoreoptionstopk) | number | <b><i>(Public Preview)</i></b> |

@@ -34,12 +36,40 @@ export interface LanguageModelCreateCoreOptions
 > This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
 > 

+Defaults to image input and English text input if not overriden.
+
 <b>Signature:</b>

 ```typescript
 expectedInputs?: LanguageModelExpected[];
 ```

+## LanguageModelCreateCoreOptions.expectedOutputs
+
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+Defaults to English text output if not overriden.
+
+<b>Signature:</b>
+
+```typescript
+expectedOutputs?: LanguageModelExpected[];
+```
+
+## LanguageModelCreateCoreOptions.monitor
+
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+The native download monitor provided by the Chrome API. This provides direct access to the monitor object and is only needed for advanced cases. For a simpler download monitoring interface, provide an `onDownloadProgress` callback to [GenerativeModel.initializeDeviceModel()](./ai.generativemodel.md#generativemodelinitializedevicemodel) instead.
+
+<b>Signature:</b>
+
+```typescript
+monitor?: (monitor: LanguageModelDownloadMonitor) => void;
+```
+
 ## LanguageModelCreateCoreOptions.temperature

 > This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
@@ -0,0 +1,55 @@
+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 %}
+
+# LanguageModelDownloadMonitor interface
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+Interface representing the Chrome Prompt API's native download monitor object.
+
+<b>Signature:</b>
+
+```typescript
+export interface LanguageModelDownloadMonitor 
+```
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [addEventListener](./ai.languagemodeldownloadmonitor.md#languagemodeldownloadmonitoraddeventlistener) | (eventType: 'downloadprogress', eventListener: (e: { loaded: number; }) =&gt; void) =&gt; void | <b><i>(Public Preview)</i></b> |
+|  [removeEventListener](./ai.languagemodeldownloadmonitor.md#languagemodeldownloadmonitorremoveeventlistener) | (eventType: 'downloadprogress', eventListener: (e: { loaded: number; }) =&gt; void) =&gt; void | <b><i>(Public Preview)</i></b> |
+
+## LanguageModelDownloadMonitor.addEventListener
+
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+<b>Signature:</b>
+
+```typescript
+addEventListener: (eventType: 'downloadprogress', eventListener: (e: {
+        loaded: number;
+    }) => void) => void;
+```
+
+## LanguageModelDownloadMonitor.removeEventListener
+
+> This API is provided as a preview for developers and may change based on feedback that we receive. Do not use this API in a production environment.
+> 
+
+<b>Signature:</b>
+
+```typescript
+removeEventListener: (eventType: 'downloadprogress', eventListener: (e: {
+        loaded: number;
+    }) => void) => void;
+```
@@ -114,6 +114,7 @@ The Firebase AI Web SDK.
 |  [InlineDataPart](./ai.inlinedatapart.md#inlinedatapart_interface) | Content part interface if the part represents an image. |
 |  [LanguageModelCreateCoreOptions](./ai.languagemodelcreatecoreoptions.md#languagemodelcreatecoreoptions_interface) | <b><i>(Public Preview)</i></b> Configures the creation of an on-device language model session. |
 |  [LanguageModelCreateOptions](./ai.languagemodelcreateoptions.md#languagemodelcreateoptions_interface) | <b><i>(Public Preview)</i></b> Configures the creation of an on-device language model session. |
+|  [LanguageModelDownloadMonitor](./ai.languagemodeldownloadmonitor.md#languagemodeldownloadmonitor_interface) | <b><i>(Public Preview)</i></b> Interface representing the Chrome Prompt API's native download monitor object. |
 |  [LanguageModelExpected](./ai.languagemodelexpected.md#languagemodelexpected_interface) | <b><i>(Public Preview)</i></b> Options for the expected inputs for an on-device language model. |
 |  [LanguageModelMessage](./ai.languagemodelmessage.md#languagemodelmessage_interface) | <b><i>(Public Preview)</i></b> An on-device language model message. |
 |  [LanguageModelMessageContent](./ai.languagemodelmessagecontent.md#languagemodelmessagecontent_interface) | <b><i>(Public Preview)</i></b> An on-device language model content object. |
@@ -19,7 +19,12 @@ import { AIError } from '../errors';
 import { expect, use } from 'chai';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
-import { chromeAdapterFactory, ChromeAdapterImpl } from './chrome-adapter';
+import {
+  chromeAdapterFactory,
+  ChromeAdapterImpl,
+  defaultExpectedInputs,
+  defaultExpectedOutputs
+} from './chrome-adapter';
 import {
  Availability,
  LanguageModel,
@@ -54,7 +59,7 @@ async function toStringArray(

 describe('ChromeAdapter', () => {
  describe('constructor', () => {
-    it('sets image as expected input type by default', async () => {
+    it('sets default expectedInputs/Outputs if createOptions not provided', async () => {
      const languageModelProvider = {
        availability: () => Promise.resolve(Availability.AVAILABLE)
      } as LanguageModel;
@@ -75,10 +80,11 @@ describe('ChromeAdapter', () => {
        ]
      });
      expect(availabilityStub).to.have.been.calledWith({
-        expectedInputs: [{ type: 'image' }]
+        expectedInputs: defaultExpectedInputs,
+        expectedOutputs: defaultExpectedOutputs
      });
    });
-    it('sets image as expected input type by default even if other onDeviceParams params are set', async () => {
+    it('sets default expectedInputs/Outputs even if other onDeviceParams params are set', async () => {
      const languageModelProvider = {
        availability: () => Promise.resolve(Availability.AVAILABLE)
      } as LanguageModel;
@@ -102,10 +108,11 @@ describe('ChromeAdapter', () => {
        ]
      });
      expect(availabilityStub).to.have.been.calledWith({
-        expectedInputs: [{ type: 'image' }]
+        expectedInputs: defaultExpectedInputs,
+        expectedOutputs: defaultExpectedOutputs
      });
    });
-    it('sets image as expected input type by default even if other createOptions params are set', async () => {
+    it('sets default expectedInputs/Outputs even if other createOptions params are set', async () => {
      const languageModelProvider = {
        availability: () => Promise.resolve(Availability.AVAILABLE)
      } as LanguageModel;
@@ -132,10 +139,11 @@ describe('ChromeAdapter', () => {
      });
      expect(availabilityStub).to.have.been.calledWith({
        topK: 22,
-        expectedInputs: [{ type: 'image' }]
+        expectedInputs: defaultExpectedInputs,
+        expectedOutputs: defaultExpectedOutputs
      });
    });
-    it('honors explicitly set expected inputs', async () => {
+    it('honors explicitly set expected inputs and outputs', async () => {
      const languageModelProvider = {
        availability: () => Promise.resolve(Availability.AVAILABLE)
      } as LanguageModel;
@@ -145,7 +153,8 @@ describe('ChromeAdapter', () => {
      ).resolves(Availability.AVAILABLE);
      const createOptions = {
        // Explicitly sets expected inputs.
-        expectedInputs: [{ type: 'text' }]
+        expectedInputs: [{ type: 'text', languages: ['en'] }],
+        expectedOutputs: [{ type: 'text', languages: ['de'] }]
      } as LanguageModelCreateOptions;
      const adapter = new ChromeAdapterImpl(
        languageModelProvider,
@@ -310,14 +319,12 @@ describe('ChromeAdapter', () => {
        })
      ).to.be.true;
    });
-    it('returns false and triggers download when model is available after download', async () => {
+    it('returns false if model is downloadable but not downloaded', async () => {
      const languageModelProvider = {
        availability: () => Promise.resolve(Availability.DOWNLOADABLE),
        create: () => Promise.resolve({})
      } as LanguageModel;
-      const createStub = stub(languageModelProvider, 'create').resolves(
-        {} as LanguageModel
-      );
+      stub(languageModelProvider, 'create').resolves({} as LanguageModel);
      const createOptions = {
        expectedInputs: [{ type: 'image' }]
      } as LanguageModelCreateOptions;
@@ -331,7 +338,51 @@ describe('ChromeAdapter', () => {
          contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
        })
      ).to.be.false;
-      expect(createStub).to.have.been.calledOnceWith(createOptions);
+    });
+    it('returns false when model is never available', async () => {
+      const languageModelProvider = {
+        availability: () => Promise.resolve(Availability.UNAVAILABLE),
+        create: () => Promise.resolve({})
+      } as LanguageModel;
+      const adapter = new ChromeAdapterImpl(
+        languageModelProvider,
+        InferenceMode.PREFER_ON_DEVICE
+      );
+      expect(
+        await adapter.isAvailable({
+          contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+        })
+      ).to.be.false;
+    });
+  });
+  describe('downloadIfAvailable', () => {
+    it('calls create() with listener if downloadable', async () => {
+      const progressCallback = stub();
+      const mockCreate = stub().resolves();
+      const adapter = new ChromeAdapterImpl(
+        {
+          availability: stub().resolves(Availability.DOWNLOADABLE),
+          create: mockCreate
+        } as unknown as LanguageModel,
+        InferenceMode.PREFER_ON_DEVICE
+      );
+      await adapter.downloadIfAvailable(progressCallback);
+      await adapter.downloadPromise;
+      expect(mockCreate.getCall(0).args[0].monitor).to.exist;
+    });
+    it('calls create() with listener if downloading', async () => {
+      const progressCallback = stub();
+      const mockCreate = stub().resolves();
+      const adapter = new ChromeAdapterImpl(
+        {
+          availability: stub().resolves(Availability.DOWNLOADING),
+          create: mockCreate
+        } as unknown as LanguageModel,
+        InferenceMode.PREFER_ON_DEVICE
+      );
+      await adapter.downloadIfAvailable(progressCallback);
+      await adapter.downloadPromise;
+      expect(mockCreate.getCall(0).args[0].monitor).to.exist;
    });
    it('avoids redundant downloads', async () => {
      const languageModelProvider = {
@@ -348,12 +399,8 @@ describe('ChromeAdapter', () => {
        languageModelProvider,
        InferenceMode.PREFER_ON_DEVICE
      );
-      await adapter.isAvailable({
-        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
-      });
-      await adapter.isAvailable({
-        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
-      });
+      await adapter.downloadIfAvailable();
+      await adapter.downloadIfAvailable();
      expect(createStub).to.have.been.calledOnce;
    });
    it('clears state when download completes', async () => {
@@ -372,30 +419,11 @@ describe('ChromeAdapter', () => {
        languageModelProvider,
        InferenceMode.PREFER_ON_DEVICE
      );
-      await adapter.isAvailable({
-        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
-      });
+      await adapter.downloadIfAvailable();
      resolveDownload!();
-      await adapter.isAvailable({
-        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
-      });
+      await adapter.downloadIfAvailable();
      expect(createStub).to.have.been.calledTwice;
    });
-    it('returns false when model is never available', async () => {
-      const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.UNAVAILABLE),
-        create: () => Promise.resolve({})
-      } as LanguageModel;
-      const adapter = new ChromeAdapterImpl(
-        languageModelProvider,
-        InferenceMode.PREFER_ON_DEVICE
-      );
-      expect(
-        await adapter.isAvailable({
-          contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
-        })
-      ).to.be.false;
-    });
  });
  describe('generateContent', () => {
    it('throws if Chrome API is undefined', async () => {
@@ -37,8 +37,20 @@ import {
  LanguageModelMessageRole
 } from '../types/language-model';

-// Defaults to support image inputs for convenience.
-const defaultExpectedInputs: LanguageModelExpected[] = [{ type: 'image' }];
+// Defaults to English text. This can be overriden by user config.
+const defaultExpectedText: LanguageModelExpected = {
+  type: 'text',
+  languages: ['en']
+};
+
+export const defaultExpectedInputs: LanguageModelExpected[] = [
+  defaultExpectedText,
+  { type: 'image' }
+];
+
+export const defaultExpectedOutputs: LanguageModelExpected[] = [
+  defaultExpectedText
+];

 /**
 * Defines an inference "backend" that uses Chrome's on-device model,
@@ -48,12 +60,13 @@ const defaultExpectedInputs: LanguageModelExpected[] = [{ type: 'image' }];
 export class ChromeAdapterImpl implements ChromeAdapter {
  // Visible for testing
  static SUPPORTED_MIME_TYPES = ['image/jpeg', 'image/png'];
-  private isDownloading = false;
-  private downloadPromise: Promise<LanguageModel | void> | undefined;
+  // Promise for on-device model download completion
+  downloadPromise: Promise<LanguageModel | void> | null = null;
  private oldSession: LanguageModel | undefined;
  onDeviceParams: OnDeviceParams = {
    createOptions: {
-      expectedInputs: defaultExpectedInputs
+      expectedInputs: defaultExpectedInputs,
+      expectedOutputs: defaultExpectedOutputs
    }
  };
  constructor(
@@ -65,11 +78,18 @@ export class ChromeAdapterImpl implements ChromeAdapter {
      this.onDeviceParams = onDeviceParams;
      if (!this.onDeviceParams.createOptions) {
        this.onDeviceParams.createOptions = {
-          expectedInputs: defaultExpectedInputs
+          expectedInputs: defaultExpectedInputs,
+          expectedOutputs: defaultExpectedOutputs
        };
-      } else if (!this.onDeviceParams.createOptions.expectedInputs) {
-        this.onDeviceParams.createOptions.expectedInputs =
-          defaultExpectedInputs;
+      } else {
+        if (!this.onDeviceParams.createOptions.expectedInputs) {
+          this.onDeviceParams.createOptions.expectedInputs =
+            defaultExpectedInputs;
+        }
+        if (!this.onDeviceParams.createOptions.expectedOutputs) {
+          this.onDeviceParams.createOptions.expectedOutputs =
+            defaultExpectedOutputs;
+        }
      }
    }
  }
@@ -81,7 +101,7 @@ export class ChromeAdapterImpl implements ChromeAdapter {
   *   the mode
   *   API existence
   *   prompt formatting
-   *   model availability, including triggering download if necessary
+   *   model availability
   *
   *
   * Pros: callers needn't be concerned with details of on-device availability.</p>
@@ -103,8 +123,9 @@ export class ChromeAdapterImpl implements ChromeAdapter {
      return false;
    }

-    // Triggers out-of-band download so model will eventually become available.
-    const availability = await this.downloadIfAvailable();
+    const availability = await this.languageModelProvider?.availability(
+      this.onDeviceParams.createOptions
+    );

    if (this.mode === InferenceMode.ONLY_ON_DEVICE) {
      // If it will never be available due to API inavailability, throw.
@@ -117,9 +138,12 @@ export class ChromeAdapterImpl implements ChromeAdapter {
        availability === Availability.DOWNLOADABLE ||
        availability === Availability.DOWNLOADING
      ) {
-        // TODO(chholland): Better user experience during download - progress?
        logger.debug(`Waiting for download of LanguageModel to complete.`);
-        await this.downloadPromise;
+        try {
+          await this.downloadPromise;
+        } catch (e) {
+          throw new AIError(AIErrorCode.ERROR, (e as Error).message);
+        }
        return true;
      }
      return true;
@@ -231,13 +255,23 @@ export class ChromeAdapterImpl implements ChromeAdapter {
  /**
   * Encapsulates logic to get availability and download a model if one is downloadable.
   */
-  private async downloadIfAvailable(): Promise<Availability | undefined> {
+  async downloadIfAvailable(
+    onDownloadProgress?: (progressValue: number) => void
+  ): Promise<Availability | undefined> {
    const availability = await this.languageModelProvider?.availability(
      this.onDeviceParams.createOptions
    );

-    if (availability === Availability.DOWNLOADABLE) {
-      this.download();
+    /**
+     * Download may have been automatically started by Chrome or other
+     * code, in which case we should still call create() to get a promise
+     * for when the download completes.
+     */
+    if (
+      availability === Availability.DOWNLOADABLE ||
+      availability === Availability.DOWNLOADING
+    ) {
+      this.download(onDownloadProgress);
    }

    return availability;
@@ -246,21 +280,28 @@ export class ChromeAdapterImpl implements ChromeAdapter {
  /**
   * Triggers out-of-band download of an on-device model.
   *
-   * Chrome only downloads models as needed. Chrome knows a model is needed when code calls
-   * LanguageModel.create.
+   * Chrome may automatically begin a download on startup or
+   * it may trigger a download when code calls LanguageModel.create.
   *
   * Since Chrome manages the download, the SDK can only avoid redundant download requests by
   * tracking if a download has previously been requested.
   */
-  private download(): void {
-    if (this.isDownloading) {
+  private download(onDownloadProgress?: (progressValue: number) => void): void {
+    if (this.downloadPromise) {
      return;
    }
-    this.isDownloading = true;
+    const options = { ...this.onDeviceParams.createOptions };
+    if (options && !options.monitor && onDownloadProgress) {
+      options.monitor = m => {
+        m.addEventListener('downloadprogress', e => {
+          onDownloadProgress(e.loaded);
+        });
+      };
+    }
    this.downloadPromise = this.languageModelProvider
-      ?.create(this.onDeviceParams.createOptions)
+      ?.create(options)
      .finally(() => {
-        this.isDownloading = false;
+        this.downloadPromise = null;
      });
  }

@@ -334,6 +375,13 @@ export class ChromeAdapterImpl implements ChromeAdapter {
        'Chrome AI requested for unsupported browser version.'
      );
    }
+    /**
+     * Note: `create()` will trigger a download as a side effect if
+     * a model is not downloaded or downloading, but this specific
+     * create() call should not be relied on for the download. The download
+     * should be triggered explicitly by the user in
+     * GenerativeModel.initializeDeviceModel().
+     */
    const newSession = await this.languageModelProvider.create(
      this.onDeviceParams.createOptions
    );
@@ -39,6 +39,7 @@ import { VertexAIBackend } from '../backend';
 import { AIError } from '../errors';
 import chaiAsPromised from 'chai-as-promised';
 import { fakeChromeAdapter } from '../../test-utils/get-fake-firebase-services';
+import { Availability } from '../types/language-model';

 use(sinonChai);
 use(chaiAsPromised);
@@ -675,8 +676,76 @@ describe('GenerativeModel', () => {
    );
  });
 });
+describe('initializeDeviceModel', () => {
+  it('throws if unavailable and ONLY_ON_DEVICE', async () => {
+    // @ts-ignore
+    const mockChromeAdapter = {
+      mode: InferenceMode.ONLY_ON_DEVICE,
+      downloadIfAvailable: stub().resolves(Availability.UNAVAILABLE),
+      download: stub()
+    };
+    const model = new GenerativeModel(
+      fakeAI,
+      { model: 'model' },
+      {},
+      //@ts-ignore
+      mockChromeAdapter
+    );
+    await expect(model.initializeDeviceModel()).to.be.rejectedWith(
+      'Local LanguageModel API not available in this environment'
+    );
+    expect(mockChromeAdapter.download).to.not.be.called;
+  });
+  it('noops if ONLY_IN_CLOUD', async () => {
+    // @ts-ignore
+    const mockChromeAdapter = {
+      mode: InferenceMode.ONLY_IN_CLOUD,
+      downloadIfAvailable: stub().resolves(Availability.AVAILABLE),
+      download: stub()
+    };
+    const model = new GenerativeModel(
+      fakeAI,
+      { model: 'model' },
+      {},
+      //@ts-ignore
+      mockChromeAdapter
+    );
+    await model.initializeDeviceModel();
+    expect(mockChromeAdapter.download).to.not.be.called;
+  });
+  it('noops if no adapter', async () => {
+    // @ts-ignore
+    const mockChromeAdapter = {
+      mode: InferenceMode.PREFER_ON_DEVICE,
+      downloadIfAvailable: stub().resolves(Availability.AVAILABLE),
+      download: stub()
+    };
+    const model = new GenerativeModel(fakeAI, { model: 'model' }, {});
+    await model.initializeDeviceModel();
+    expect(mockChromeAdapter.download).to.not.be.called;
+  });
+  it('passes downloadProgress callback to download()', async () => {
+    // @ts-ignore
+    const mockChromeAdapter = {
+      mode: InferenceMode.PREFER_ON_DEVICE,
+      downloadIfAvailable: stub().resolves(Availability.AVAILABLE)
+    };
+    const model = new GenerativeModel(
+      fakeAI,
+      { model: 'model' },
+      {},
+      //@ts-ignore
+      mockChromeAdapter
+    );
+    const progressCallback = (): void => {};
+    await model.initializeDeviceModel(progressCallback);
+    expect(mockChromeAdapter.downloadIfAvailable).to.be.calledWith(
+      progressCallback
+    );
+  });
+});

-describe('GenerativeModel dispatch logic', () => {
+describe('GenerativeModel hybrid dispatch logic', () => {
  let makeRequestStub: SinonStub;
  let mockChromeAdapter: ChromeAdapter;

@@ -42,10 +42,13 @@ import {
  formatGenerateContentInput,
  formatSystemInstruction
 } from '../requests/request-helpers';
-import { AI, AIErrorCode } from '../public-types';
+import { AI, AIErrorCode, InferenceMode } from '../public-types';
 import { AIModel } from './ai-model';
 import { ChromeAdapter } from '../types/chrome-adapter';
 import { AIError } from '../errors';
+import { ChromeAdapterImpl } from '../methods/chrome-adapter';
+import { Availability } from '../types/language-model';
+import { logger } from '../logger';

 /**
 * Class for generative model APIs.
@@ -77,6 +80,62 @@ export class GenerativeModel extends AIModel {
    this.requestOptions = requestOptions || {};
  }

+  /**
+   * Initializes on-device models.
+   *
+   * @remarks
+   * This may trigger a download on first
+   * use. Wait for this promise to complete before calling inference
+   * methods if you want to ensure the device models are ready before
+   * any calls. Calling inference methods before the device is ready
+   * will result in a cloud fallback if `inferenceMode` is set to
+   * `PREFER_ON_DEVICE`, and an error if set to `ONLY_ON_DEVICE`.
+   *
+   * IMPORTANT: This call must be made on or after a user has interacted
+   * with the page (for example, through a button click or key press).
+   * If it is called without a user interaction, and it requires a download,
+   * this will cause an error.
+   *
+   * See the
+   * {@link https://developer.chrome.com/docs/ai/prompt-api#use_the_prompt_api | Prompt API docs }
+   * for more details on this requirement.
+   *
+   * @param onDownloadProgress A callback called repeatedly as the
+   * download progresses that provides a `progressValue` between 0
+   * and 1 representing how much of the download is complete. This
+   * will be ignored if `monitor` was populated in
+   * {@link LanguageModelCreateOptions}.
+   *
+   * @public
+   */
+
+  async initializeDeviceModel(
+    onDownloadProgress?: (progressValue: number) => void
+  ): Promise<void> {
+    if (
+      !this.chromeAdapter ||
+      this.chromeAdapter.mode === InferenceMode.ONLY_IN_CLOUD
+    ) {
+      return;
+    }
+    const availability = await (
+      this.chromeAdapter as ChromeAdapterImpl
+    ).downloadIfAvailable(onDownloadProgress);
+    if (availability === Availability.UNAVAILABLE) {
+      const notEnabledError = new AIError(
+        AIErrorCode.API_NOT_ENABLED,
+        'Local LanguageModel API not available in this environment.'
+      );
+      if (this.chromeAdapter.mode === InferenceMode.ONLY_ON_DEVICE) {
+        throw notEnabledError;
+      } else {
+        // No reason to throw if not in ONLY_ON_DEVICE mode.
+        logger.debug(notEnabledError.message);
+      }
+    }
+    await (this.chromeAdapter as ChromeAdapterImpl).downloadPromise;
+  }
+
  /**
   * Makes a single non-streaming call to the model
   * and returns an object containing a single {@link GenerateContentResponse}.
@@ -42,7 +42,8 @@ describe('callCloudOrDevice', () => {
    // @ts-ignore
    chromeAdapter = {
      mode: InferenceMode.PREFER_ON_DEVICE,
-      isAvailable: stub(),
+      // @ts-ignore
+      isAvailable: stub().resolves(true),
      generateContent: stub(),
      generateContentStream: stub(),
      countTokens: stub()
@@ -84,7 +84,11 @@ export async function callCloudOrDevice<Response>(
          inferenceSource: InferenceSource.IN_CLOUD
        };
      } catch (e) {
-        if (e instanceof AIError && errorsCausingFallback.includes(e.code)) {
+        if (
+          e instanceof AIError &&
+          errorsCausingFallback.includes(e.code) &&
+          (await chromeAdapter.isAvailable(request))
+        ) {
          return {
            response: await onDeviceCall(),
            inferenceSource: InferenceSource.ON_DEVICE
@@ -32,6 +32,7 @@ export {
  LanguageModelMessageContentValue,
  LanguageModelMessageRole,
  LanguageModelMessageType,
-  LanguageModelPromptOptions
+  LanguageModelPromptOptions,
+  LanguageModelDownloadMonitor
 } from './language-model';
 export * from './chrome-adapter';
@@ -62,7 +62,40 @@ export interface LanguageModelCreateCoreOptions {
   * @deprecated
   */
  temperature?: number;
+  /**
+   * Defaults to image input and English text input
+   * if not overriden.
+   */
  expectedInputs?: LanguageModelExpected[];
+  /**
+   * Defaults to English text output if not overriden.
+   */
+  expectedOutputs?: LanguageModelExpected[];
+  /**
+   * The native download monitor provided by the Chrome API.
+   * This provides direct access to the monitor object and is
+   * only needed for advanced cases. For
+   * a simpler download monitoring interface, provide an
+   * `onDownloadProgress` callback
+   * to {@link GenerativeModel.initializeDeviceModel} instead.
+   */
+  monitor?: (monitor: LanguageModelDownloadMonitor) => void;
+}
+
+/**
+ * Interface representing the Chrome Prompt API's native
+ * download monitor object.
+ * @beta
+ */
+export interface LanguageModelDownloadMonitor {
+  addEventListener: (
+    eventType: 'downloadprogress',
+    eventListener: (e: { loaded: number }) => void
+  ) => void;
+  removeEventListener: (
+    eventType: 'downloadprogress',
+    eventListener: (e: { loaded: number }) => void
+  ) => void;
 }

 /**