diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 550b87c9e0b..d7b73f9a256 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -53,7 +53,7 @@ async function toStringArray( describe('ChromeAdapter', () => { describe('constructor', () => { - it('sets image as expected input type by default', async () => { + it('determines expected inputs by request inspection', async () => { const languageModelProvider = { availability: () => Promise.resolve(Availability.available) } as LanguageModel; @@ -69,7 +69,11 @@ describe('ChromeAdapter', () => { contents: [ { role: 'user', - parts: [{ text: 'hi' }] + parts: [ + { text: 'hi' }, + // Triggers image as expected type. + { inlineData: { mimeType: 'image/asd', data: 'asd' } } + ] } ] }); diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index 9ac8f350a02..76575ff0e07 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -27,7 +27,9 @@ import { Availability, LanguageModel, LanguageModelCreateOptions, - LanguageModelMessageContent + LanguageModelExpectedInput, + LanguageModelMessageContent, + LanguageModelMessageType } from '../types/language-model'; /** @@ -44,9 +46,7 @@ export class ChromeAdapter { private languageModelProvider?: LanguageModel, private mode?: InferenceMode, private onDeviceParams: LanguageModelCreateOptions = {} - ) { - this.addImageTypeAsExpectedInput(); - } + ) {} /** * Checks if a given request can be made on-device. @@ -68,8 +68,10 @@ export class ChromeAdapter { return false; } + const expectedInputs = ChromeAdapter.extractExpectedInputs(request); + // Triggers out-of-band download so model will eventually become available. - const availability = await this.downloadIfAvailable(); + const availability = await this.downloadIfAvailable(expectedInputs); if (this.mode === 'only_on_device') { return true; @@ -129,6 +131,33 @@ export class ChromeAdapter { ); } + /** + * Maps + * + * Vertex's input mime types to + * + * Chrome's expected types. + * + *
Chrome's API checks availability by type. It's tedious to specify the types in advance, so + * this method infers the types.
+ */ + private static extractExpectedInputs( + request: GenerateContentRequest + ): LanguageModelExpectedInput[] { + const inputSet = new Set