adding image as input to create language model object

Author: gsiddh

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

@@ -173,18 +173,19 @@ describe('ChromeAdapter', () => {
       const createStub = stub(languageModelProvider, 'create').resolves(
         {} as LanguageModel
       );
-      const onDeviceParams = {} as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
-        'prefer_on_device',
-        onDeviceParams
+        'prefer_on_device'
       );
+      const expectedOnDeviceParams = {
+        expectedInputs: [{ type: 'image' }]
+      } as LanguageModelCreateOptions;
       expect(
         await adapter.isAvailable({
           contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
         })
       ).to.be.false;
-      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
+      expect(createStub).to.have.been.calledOnceWith(expectedOnDeviceParams);
     });
     it('avoids redundant downloads', async () => {
       const languageModelProvider = {
@@ -261,6 +262,10 @@ describe('ChromeAdapter', () => {
       const onDeviceParams = {
         systemPrompt: 'be yourself'
       } as LanguageModelCreateOptions;
+      const expectedOnDeviceParams = {
+        systemPrompt: 'be yourself',
+        expectedInputs: [{ type: 'image' }]
+      } as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
         'prefer_on_device',
@@ -271,7 +276,7 @@ describe('ChromeAdapter', () => {
       } as GenerateContentRequest;
       const response = await adapter.generateContent(request);
       // Asserts initialization params are proxied.
-      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
+      expect(createStub).to.have.been.calledOnceWith(expectedOnDeviceParams);
       // Asserts Vertex input type is mapped to Chrome type.
       expect(promptStub).to.have.been.calledOnceWith([
         {
@@ -306,6 +311,10 @@ describe('ChromeAdapter', () => {
       const onDeviceParams = {
         systemPrompt: 'be yourself'
       } as LanguageModelCreateOptions;
+      const expectedOnDeviceParams = {
+        systemPrompt: 'be yourself',
+        expectedInputs: [{ type: 'image' }]
+      } as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
         'prefer_on_device',
@@ -329,7 +338,7 @@ describe('ChromeAdapter', () => {
       } as GenerateContentRequest;
       const response = await adapter.generateContent(request);
       // Asserts initialization params are proxied.
-      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
+      expect(createStub).to.have.been.calledOnceWith(expectedOnDeviceParams);
       // Asserts Vertex input type is mapped to Chrome type.
       expect(promptStub).to.have.been.calledOnceWith([
         {
@@ -360,6 +369,10 @@ describe('ChromeAdapter', () => {
       const onDeviceParams = {
         systemPrompt: 'be yourself'
       } as LanguageModelCreateOptions;
+      const expectedOnDeviceParams = {
+        systemPrompt: 'be yourself',
+        expectedInputs: [{ type: 'image' }]
+      } as LanguageModelCreateOptions;
 
       // setting up stubs
       const languageModelProvider = {
@@ -389,7 +402,7 @@ describe('ChromeAdapter', () => {
       } as GenerateContentRequest;
       const response = await adapter.countTokens(countTokenRequest);
       // Asserts initialization params are proxied.
-      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
+      expect(createStub).to.have.been.calledOnceWith(expectedOnDeviceParams);
       // Asserts Vertex input type is mapped to Chrome type.
       expect(measureInputUsageStub).to.have.been.calledOnceWith([
         {
@@ -423,6 +436,9 @@ describe('ChromeAdapter', () => {
         })
       );
       const onDeviceParams = {} as LanguageModelCreateOptions;
+      const expectedOnDeviceParams = {
+        expectedInputs: [{ type: 'image' }]
+      } as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
         'prefer_on_device',
@@ -432,7 +448,7 @@ describe('ChromeAdapter', () => {
         contents: [{ role: 'user', parts: [{ text: 'anything' }] }]
       } as GenerateContentRequest;
       const response = await adapter.generateContentStream(request);
-      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
+      expect(createStub).to.have.been.calledOnceWith(expectedOnDeviceParams);
       expect(promptStub).to.have.been.calledOnceWith([
         {
           type: 'text',
@@ -464,6 +480,9 @@ describe('ChromeAdapter', () => {
         })
       );
       const onDeviceParams = {} as LanguageModelCreateOptions;
+      const expectedOnDeviceParams = {
+        expectedInputs: [{ type: 'image' }]
+      } as LanguageModelCreateOptions;
       const adapter = new ChromeAdapter(
         languageModelProvider,
         'prefer_on_device',
@@ -486,7 +505,7 @@ describe('ChromeAdapter', () => {
         ]
       } as GenerateContentRequest;
       const response = await adapter.generateContentStream(request);
-      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
+      expect(createStub).to.have.been.calledOnceWith(expectedOnDeviceParams);
       expect(promptStub).to.have.been.calledOnceWith([
         {
           type: 'text',

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

@@ -162,10 +162,6 @@ export class ChromeAdapter {
       if (content.role === 'function') {
         return false;
       }
-
-      if (!content.parts[0].text) {
-        return false;
-      }
     }
 
     return true;
@@ -185,8 +181,10 @@ export class ChromeAdapter {
       return;
     }
     this.isDownloading = true;
+    const options = this.onDeviceParams || {};
+    ChromeAdapter.addImageTypeAsExpectedInput(options);
     this.downloadPromise = this.languageModelProvider
-      ?.create(this.onDeviceParams)
+      ?.create(options)
       .then(() => {
         this.isDownloading = false;
       });
@@ -204,7 +202,6 @@ export class ChromeAdapter {
         content: part.text
       };
     } else if (part.inlineData) {
-      // this is for the image type
       const formattedImageContent = await fetch(
         `data:${part.inlineData.mimeType};base64,${part.inlineData.data}`
       );
@@ -235,6 +232,7 @@ export class ChromeAdapter {
     options: LanguageModelCreateOptions
   ): Promise<LanguageModel> {
     // TODO: could we use this.onDeviceParams instead of passing in options?
+    ChromeAdapter.addImageTypeAsExpectedInput(options);
     const newSession = await this.languageModelProvider!.create(options);
     if (this.oldSession) {
       this.oldSession.destroy();
@@ -244,6 +242,13 @@ export class ChromeAdapter {
     return newSession;
   }
 
+  private static addImageTypeAsExpectedInput(
+    options: LanguageModelCreateOptions
+  ): void {
+    options.expectedInputs = options.expectedInputs || [];
+    options.expectedInputs.push({ type: 'image' });
+  }
+
   /**
    * Formats string returned by Chrome as a {@link Response} returned by Vertex.
    */