Merge 5ee53bf into dc3794a

Author: erikeldridge

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

@@ -52,6 +52,59 @@ async function toStringArray(
 }
 
 describe('ChromeAdapter', () => {
+  describe('constructor', () => {
+    it('sets image as expected input type by default', async () => {
+      const languageModelProvider = {
+        availability: () => Promise.resolve(Availability.available)
+      } as LanguageModel;
+      const availabilityStub = stub(
+        languageModelProvider,
+        'availability'
+      ).resolves(Availability.available);
+      const adapter = new ChromeAdapter(
+        languageModelProvider,
+        'prefer_on_device'
+      );
+      await adapter.isAvailable({
+        contents: [
+          {
+            role: 'user',
+            parts: [{ text: 'hi' }]
+          }
+        ]
+      });
+      expect(availabilityStub).to.have.been.calledWith({
+        expectedInputs: [{ type: 'image' }]
+      });
+    });
+    it('honors explicitly set expected inputs', async () => {
+      const languageModelProvider = {
+        availability: () => Promise.resolve(Availability.available)
+      } as LanguageModel;
+      const availabilityStub = stub(
+        languageModelProvider,
+        'availability'
+      ).resolves(Availability.available);
+      const onDeviceParams = {
+        // Explicitly sets expected inputs.
+        expectedInputs: [{ type: 'text' }]
+      } as LanguageModelCreateOptions;
+      const adapter = new ChromeAdapter(
+        languageModelProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
+      await adapter.isAvailable({
+        contents: [
+          {
+            role: 'user',
+            parts: [{ text: 'hi' }]
+          }
+        ]
+      });
+      expect(availabilityStub).to.have.been.calledWith(onDeviceParams);
+    });
+  });
   describe('isAvailable', () => {
     it('returns false if mode is only cloud', async () => {
       const adapter = new ChromeAdapter(undefined, 'only_in_cloud');
@@ -100,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)
@@ -110,7 +190,15 @@ describe('ChromeAdapter', () => {
       );
       expect(
         await adapter.isAvailable({
-          contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
+          contents: [
+            {
+              role: 'user',
+              parts: [
+                { text: 'describe this image' },
+                { inlineData: { mimeType: 'image/jpeg', data: 'asd' } }
+              ]
+            }
+          ]
         })
       ).to.be.true;
     });

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

@@ -35,14 +35,18 @@ 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'];
   private isDownloading = false;
   private downloadPromise: Promise<LanguageModel | void> | undefined;
   private oldSession: LanguageModel | undefined;
   constructor(
     private languageModelProvider?: LanguageModel,
     private mode?: InferenceMode,
     private onDeviceParams: LanguageModelCreateOptions = {}
-  ) {}
+  ) {
+    this.addImageTypeAsExpectedInput();
+  }
 
   /**
    * Checks if a given request can be made on-device.
@@ -140,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;
@@ -236,6 +252,11 @@ export class ChromeAdapter {
     return newSession;
   }
 
+  private addImageTypeAsExpectedInput(): void {
+    // Defaults to support image inputs for convenience.
+    this.onDeviceParams.expectedInputs ??= [{ type: 'image' }];
+  }
+
   /**
    * Formats string returned by Chrome as a {@link Response} returned by Vertex.
    */