diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index abdbd08c401..550b87c9e0b 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/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; }); diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index 521e9ca7101..9ac8f350a02 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/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']; private isDownloading = false; private downloadPromise: Promise | undefined; private oldSession: LanguageModel | undefined; @@ -42,7 +44,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. @@ -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. */