Support image input by default

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');
@@ -110,7 +163,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

@@ -42,7 +42,9 @@ export class ChromeAdapter {
     private languageModelProvider?: LanguageModel,
     private mode?: InferenceMode,
     private onDeviceParams: LanguageModelCreateOptions = {}
-  ) {}
+  ) {
+    this.addImageTypeAsExpectedInput();
+  }
 
   /**
    * Checks if a given request can be made on-device.
@@ -236,6 +238,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.
    */