VinF Hybrid Inference: throw if only_on_device and model is unavailable (#8965)

Author: erikeldridge

e2e/sample-apps/modular.js

@@ -314,7 +314,7 @@ async function callVertexAI(app) {
   console.log('[VERTEXAI] start');
   const vertexAI = getVertexAI(app);
   const model = getGenerativeModel(vertexAI, {
-    mode: 'prefer_on_device'
+    mode: 'only_on_device'
   });
   const singleResult = await model.generateContent([
     { text: 'describe the following:' },

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

@@ -61,19 +61,8 @@ describe('ChromeAdapter', () => {
         })
       ).to.be.false;
     });
-    it('returns false if AI API is undefined', async () => {
-      const adapter = new ChromeAdapter(undefined, 'prefer_on_device');
-      expect(
-        await adapter.isAvailable({
-          contents: []
-        })
-      ).to.be.false;
-    });
     it('returns false if LanguageModel API is undefined', async () => {
-      const adapter = new ChromeAdapter(
-        {} as LanguageModel,
-        'prefer_on_device'
-      );
+      const adapter = new ChromeAdapter(undefined, 'prefer_on_device');
       expect(
         await adapter.isAvailable({
           contents: []
@@ -82,7 +71,9 @@ describe('ChromeAdapter', () => {
     });
     it('returns false if request contents empty', async () => {
       const adapter = new ChromeAdapter(
-        {} as LanguageModel,
+        {
+          availability: async () => Availability.available
+        } as LanguageModel,
         'prefer_on_device'
       );
       expect(
@@ -93,7 +84,9 @@ describe('ChromeAdapter', () => {
     });
     it('returns false if request content has function role', async () => {
       const adapter = new ChromeAdapter(
-        {} as LanguageModel,
+        {
+          availability: async () => Availability.available
+        } as LanguageModel,
         'prefer_on_device'
       );
       expect(
@@ -107,51 +100,6 @@ describe('ChromeAdapter', () => {
         })
       ).to.be.false;
     });
-    it('returns false if request system instruction has function role', async () => {
-      const adapter = new ChromeAdapter(
-        {} as LanguageModel,
-        'prefer_on_device'
-      );
-      expect(
-        await adapter.isAvailable({
-          contents: [],
-          systemInstruction: {
-            role: 'function',
-            parts: []
-          }
-        })
-      ).to.be.false;
-    });
-    it('returns false if request system instruction has multiple parts', async () => {
-      const adapter = new ChromeAdapter(
-        {} as LanguageModel,
-        'prefer_on_device'
-      );
-      expect(
-        await adapter.isAvailable({
-          contents: [],
-          systemInstruction: {
-            role: 'function',
-            parts: [{ text: 'a' }, { text: 'b' }]
-          }
-        })
-      ).to.be.false;
-    });
-    it('returns false if request system instruction has non-text part', async () => {
-      const adapter = new ChromeAdapter(
-        {} as LanguageModel,
-        'prefer_on_device'
-      );
-      expect(
-        await adapter.isAvailable({
-          contents: [],
-          systemInstruction: {
-            role: 'function',
-            parts: [{ inlineData: { mimeType: 'a', data: 'b' } }]
-          }
-        })
-      ).to.be.false;
-    });
     it('returns true if model is readily available', async () => {
       const languageModelProvider = {
         availability: () => Promise.resolve(Availability.available)
@@ -246,7 +194,20 @@ describe('ChromeAdapter', () => {
       ).to.be.false;
     });
   });
-  describe('generateContentOnDevice', () => {
+  describe('generateContent', () => {
+    it('throws if Chrome API is undefined', async () => {
+      const adapter = new ChromeAdapter(undefined, 'only_on_device');
+      await expect(
+        adapter.generateContent({
+          contents: []
+        })
+      )
+        .to.eventually.be.rejectedWith(
+          VertexAIError,
+          'Chrome AI requested for unsupported browser version.'
+        )
+        .and.have.property('code', VertexAIErrorCode.REQUEST_ERROR);
+    });
     it('generates content', async () => {
       const languageModelProvider = {
         create: () => Promise.resolve({})

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

@@ -60,29 +60,26 @@ export class ChromeAdapter {
    * separation of concerns.</p>
    */
   async isAvailable(request: GenerateContentRequest): Promise<boolean> {
-    // Returns false if we should only use in-cloud inference.
     if (this.mode === 'only_in_cloud') {
       return false;
     }
-    // Returns false if the on-device inference API is undefined.;
-    if (!this.languageModelProvider) {
-      return false;
-    }
-    // Returns false if the request can't be run on-device.
-    if (!ChromeAdapter.isOnDeviceRequest(request)) {
-      return false;
+
+    const availability = await this.languageModelProvider?.availability();
+
+    // Triggers async model download so it'll be available next time.
+    if (availability === Availability.downloadable) {
+      this.download();
     }
-    const availability = await this.languageModelProvider.availability();
-    switch (availability) {
-      case Availability.available:
-        // Returns true only if a model is immediately available.
-        return true;
-      case Availability.downloadable:
-        // Triggers async download if model is downloadable.
-        this.download();
-      default:
-        return false;
+
+    if (this.mode === 'only_on_device') {
+      return true;
     }
+
+    // Applies prefer_on_device logic.
+    return (
+      availability === Availability.available &&
+      ChromeAdapter.isOnDeviceRequest(request)
+    );
   }
 
   /**
@@ -221,6 +218,12 @@ export class ChromeAdapter {
     // TODO: define a default value, since these are optional.
     options: LanguageModelCreateOptions
   ): Promise<LanguageModel> {
+    if (!this.languageModelProvider) {
+      throw new VertexAIError(
+        VertexAIErrorCode.REQUEST_ERROR,
+        'Chrome AI requested for unsupported browser version.'
+      );
+    }
     // TODO: could we use this.onDeviceParams instead of passing in options?
     ChromeAdapter.addImageTypeAsExpectedInput(options);
     const newSession = await this.languageModelProvider!.create(options);