Guard against unsupported image types

Author: erikeldridge

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

@@ -153,6 +153,33 @@ describe('ChromeAdapter', () => {
         })
       ).to.be.false;
     });
+    it('returns true if request has image with supported mime type', async () => {
+      const adapter = new ChromeAdapter(
+        {
+          availability: async () => Availability.available
+        } as LanguageModel,
+        'prefer_on_device'
+      );
+      for (const mimeType of ChromeAdapter.SUPPORTED_MIME_TYPES) {
+        expect(
+          await adapter.isAvailable({
+            contents: [
+              {
+                role: 'user',
+                parts: [
+                  {
+                    inlineData: {
+                      mimeType,
+                      data: ''
+                    }
+                  }
+                ]
+              }
+            ]
+          })
+        ).to.be.true;
+      }
+    });
     it('returns true if model is readily available', async () => {
       const languageModelProvider = {
         availability: () => Promise.resolve(Availability.available)

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

@@ -35,6 +35,8 @@ import {
  * and encapsulates logic for detecting when on-device is possible.
  */
 export class ChromeAdapter {
+  // Visible for testing
+  static SUPPORTED_MIME_TYPES = ['image/jpeg', 'image/png', 'image/webp'];
   private isDownloading = false;
   private downloadPromise: Promise<LanguageModel | void> | undefined;
   private oldSession: LanguageModel | undefined;
@@ -142,6 +144,18 @@ export class ChromeAdapter {
       if (content.role !== 'user') {
         return false;
       }
+
+      // Returns false if request contains an image with an unsupported mime type.
+      for (const part of content.parts) {
+        if (
+          part.inlineData &&
+          ChromeAdapter.SUPPORTED_MIME_TYPES.indexOf(
+            part.inlineData.mimeType
+          ) === -1
+        ) {
+          return false;
+        }
+      }
     }
 
     return true;