Test model download logic

Author: erikeldridge

packages/vertexai/src/methods/chrome-adapter.test.ts

@@ -20,6 +20,7 @@ import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import { ChromeAdapter } from './chrome-adapter';
 import { Availability, LanguageModel } from '../types/language-model';
+import { stub } from 'sinon';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -56,5 +57,81 @@ describe('ChromeAdapter', () => {
         })
       ).to.be.true;
     });
+    it('returns false and triggers download when model is available after download', async () => {
+      const languageModelProvider = {
+        availability: () => Promise.resolve(Availability.downloadable),
+        create: () => Promise.resolve({})
+      } as LanguageModel;
+      const createStub = stub(languageModelProvider, 'create').resolves(
+        {} as LanguageModel
+      );
+      const adapter = new ChromeAdapter(
+        languageModelProvider,
+        'prefer_on_device'
+      );
+      expect(
+        await adapter.isAvailable({
+          contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+        })
+      ).to.be.false;
+      expect(createStub).to.have.been.calledOnce;
+    });
+    it('avoids redundant downloads', async () => {
+      const languageModelProvider = {
+        availability: () => Promise.resolve(Availability.downloadable),
+        create: () => Promise.resolve({})
+      } as LanguageModel;
+      const downloadPromise = new Promise<LanguageModel>(() => {
+        /* never resolves */
+      });
+      const createStub = stub(languageModelProvider, 'create').returns(
+        downloadPromise
+      );
+      const adapter = new ChromeAdapter(languageModelProvider);
+      await adapter.isAvailable({
+        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+      });
+      await adapter.isAvailable({
+        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+      });
+      expect(createStub).to.have.been.calledOnce;
+    });
+    it('clears state when download completes', async () => {
+      const languageModelProvider = {
+        availability: () => Promise.resolve(Availability.downloadable),
+        create: () => Promise.resolve({})
+      } as LanguageModel;
+      let resolveDownload;
+      const downloadPromise = new Promise<LanguageModel>(resolveCallback => {
+        resolveDownload = resolveCallback;
+      });
+      const createStub = stub(languageModelProvider, 'create').returns(
+        downloadPromise
+      );
+      const adapter = new ChromeAdapter(languageModelProvider);
+      await adapter.isAvailable({
+        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+      });
+      resolveDownload!();
+      await adapter.isAvailable({
+        contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+      });
+      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 ChromeAdapter(
+        languageModelProvider,
+        'prefer_on_device'
+      );
+      expect(
+        await adapter.isAvailable({
+          contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+        })
+      ).to.be.false;
+    });
   });
 });

packages/vertexai/src/methods/chrome-adapter.ts

@@ -37,8 +37,9 @@ import { isChrome } from '@firebase/util';
  * and encapsulates logic for detecting when on-device is possible.
  */
 export class ChromeAdapter {
-  downloadPromise: Promise<LanguageModel> | undefined;
-  oldSession: LanguageModel | undefined;
+  private isDownloading = false;
+  private downloadPromise: Promise<LanguageModel | void> | undefined;
+  private oldSession: LanguageModel | undefined;
   constructor(
     private languageModelProvider?: LanguageModel,
     private mode?: InferenceMode,
@@ -152,16 +153,15 @@ export class ChromeAdapter {
     return true;
   }
   private download(): void {
-    if (this.downloadPromise) {
+    if (this.isDownloading) {
       return;
     }
+    this.isDownloading = true;
     this.downloadPromise = this.languageModelProvider
       ?.create(this.onDeviceParams)
-      .then((model: LanguageModel) => {
-        delete this.downloadPromise;
-        return model;
+      .then(() => {
+        this.isDownloading = false;
       });
-    return;
   }
   private static toSystemPrompt(
     prompt: string | Content | Part | undefined