From da74e3992165be66951e12be424d84fdc006afd8 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Wed, 16 Apr 2025 10:43:07 -0700 Subject: [PATCH 1/6] Inject LanguageModel provider --- packages/vertexai/src/api.ts | 7 ++++++- packages/vertexai/src/methods/chrome-adapter.ts | 6 +++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/packages/vertexai/src/api.ts b/packages/vertexai/src/api.ts index 236ca73ce87..2f6de198608 100644 --- a/packages/vertexai/src/api.ts +++ b/packages/vertexai/src/api.ts @@ -31,6 +31,7 @@ import { import { VertexAIError } from './errors'; import { VertexAIModel, GenerativeModel, ImagenModel } from './models'; import { ChromeAdapter } from './methods/chrome-adapter'; +import { LanguageModel } from './types/language-model'; export { ChatSession } from './methods/chat-session'; export * from './requests/schema-builder'; @@ -95,7 +96,11 @@ export function getGenerativeModel( return new GenerativeModel( vertexAI, inCloudParams, - new ChromeAdapter(hybridParams.mode, hybridParams.onDeviceParams), + new ChromeAdapter( + window.LanguageModel as LanguageModel, + hybridParams.mode, + hybridParams.onDeviceParams + ), requestOptions ); } diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index 26ecd55c2da..a51b4060e26 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -16,7 +16,10 @@ */ import { GenerateContentRequest, InferenceMode } from '../types'; -import { LanguageModelCreateOptions } from '../types/language-model'; +import { + LanguageModel, + LanguageModelCreateOptions +} from '../types/language-model'; /** * Defines an inference "backend" that uses Chrome's on-device model, @@ -24,6 +27,7 @@ import { LanguageModelCreateOptions } from '../types/language-model'; */ export class ChromeAdapter { constructor( + private languageModelProvider?: LanguageModel, private mode?: InferenceMode, private onDeviceParams?: LanguageModelCreateOptions ) {} From be82cb1c7a17d7b3d6053159a382e8c8df16a297 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Tue, 1 Apr 2025 17:14:19 -0700 Subject: [PATCH 2/6] Implement ChromeAdapter class --- common/api-review/util.api.md | 5 + packages/util/src/environment.ts | 6 + .../src/methods/chrome-adapter.test.ts | 60 ++++++ .../vertexai/src/methods/chrome-adapter.ts | 173 +++++++++++++++++- packages/vertexai/src/types/language-model.ts | 6 +- 5 files changed, 241 insertions(+), 9 deletions(-) create mode 100644 packages/vertexai/src/methods/chrome-adapter.test.ts diff --git a/common/api-review/util.api.md b/common/api-review/util.api.md index 8c62ff229ac..fb8afb2319c 100644 --- a/common/api-review/util.api.md +++ b/common/api-review/util.api.md @@ -264,6 +264,11 @@ export function isBrowser(): boolean; // @public (undocumented) export function isBrowserExtension(): boolean; +// Warning: (ae-missing-release-tag) "isChrome" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal) +// +// @public (undocumented) +export function isChrome(): boolean; + // Warning: (ae-missing-release-tag) "isCloudflareWorker" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal) // // @public diff --git a/packages/util/src/environment.ts b/packages/util/src/environment.ts index a0467b08c59..50d5f534106 100644 --- a/packages/util/src/environment.ts +++ b/packages/util/src/environment.ts @@ -173,6 +173,12 @@ export function isSafari(): boolean { ); } +export function isChrome(): boolean { + return ( + !isNode() && !!navigator.userAgent && navigator.userAgent.includes('Chrome') + ); +} + /** * This method checks if indexedDB is supported by current browser/service worker context * @return true if indexedDB is supported by current browser/service worker context diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts new file mode 100644 index 00000000000..453720c121d --- /dev/null +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -0,0 +1,60 @@ +/** + * @license + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { expect, use } from 'chai'; +import sinonChai from 'sinon-chai'; +import chaiAsPromised from 'chai-as-promised'; +import { ChromeAdapter } from './chrome-adapter'; +import { Availability, LanguageModel } from '../types/language-model'; + +use(sinonChai); +use(chaiAsPromised); + +describe('ChromeAdapter', () => { + describe('isOnDeviceRequest', () => { + it('returns true for simple text part', async () => { + expect( + ChromeAdapter._isOnDeviceRequest({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }) + ).to.be.true; + }); + it('returns false if contents empty', async () => { + expect( + ChromeAdapter._isOnDeviceRequest({ + contents: [] + }) + ).to.be.false; + }); + }); + describe('isAvailable', () => { + it('returns true if a model is available', async () => { + const languageModelProvider = { + availability: () => Promise.resolve(Availability.available) + } as LanguageModel; + const adapter = new ChromeAdapter( + languageModelProvider, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }) + ).to.be.true; + }); + }); +}); diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index a51b4060e26..e5364971479 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -15,41 +15,202 @@ * limitations under the License. */ -import { GenerateContentRequest, InferenceMode } from '../types'; import { + Content, + GenerateContentRequest, + InferenceMode, + Part, + Role, + TextPart +} from '../types'; +import { + Availability, LanguageModel, - LanguageModelCreateOptions + LanguageModelCreateOptions, + LanguageModelMessageRole, + LanguageModelMessageShorthand } from '../types/language-model'; +import { isChrome } from '@firebase/util'; /** * Defines an inference "backend" that uses Chrome's on-device model, * and encapsulates logic for detecting when on-device is possible. */ export class ChromeAdapter { + downloadPromise: Promise | undefined; + oldSession: LanguageModel | undefined; constructor( private languageModelProvider?: LanguageModel, private mode?: InferenceMode, private onDeviceParams?: LanguageModelCreateOptions ) {} - // eslint-disable-next-line @typescript-eslint/no-unused-vars + /** + * Convenience method to check if a given request can be made on-device. + * Encapsulates a few concerns: 1) the mode, 2) API existence, 3) prompt formatting, and + * 4) model availability, including triggering download if necessary. + * Pros: caller needn't be concerned with details of on-device availability. Cons: this method + * spans a few concerns and splits request validation from usage. If instance variables weren't + * already part of the API, we could consider a better separation of concerns. + */ async isAvailable(request: GenerateContentRequest): Promise { - return false; + // Returns false if we should only use in-cloud inference. + if (this.mode === 'only_in_cloud') { + return false; + } + // Returns false because only Chrome's experimental Prompt API is supported. + if (!isChrome()) { + return false; + } + // Returns false if the on-device inference API is undefined.; + if (!this.languageModelProvider) { + return false; + } + // Returns false if the request can't be run on-device. + if (!ChromeAdapter._isOnDeviceRequest(request)) { + return false; + } + const availability = await this.languageModelProvider.availability(); + switch (availability) { + case Availability.available: + // Returns true only if a model is immediately available. + return true; + case Availability.downloadable: + // Triggers async download if model is downloadable. + this.download(); + default: + return false; + } } async generateContentOnDevice( - // eslint-disable-next-line @typescript-eslint/no-unused-vars request: GenerateContentRequest ): Promise { + const initialPrompts = ChromeAdapter.toInitialPrompts(request.contents); + // Assumes validation asserted there is at least one initial prompt. + const prompt = initialPrompts.pop()!; + const systemPrompt = ChromeAdapter.toSystemPrompt( + request.systemInstruction + ); + const session = await this.session({ + initialPrompts, + systemPrompt + }); + const text = await session.prompt(prompt.content); return { json: () => Promise.resolve({ candidates: [ { content: { - parts: [{ text: '' }] + parts: [{ text }] } } ] }) } as Response; } + // Visible for testing + static _isOnDeviceRequest(request: GenerateContentRequest): boolean { + if (request.systemInstruction) { + const systemContent = request.systemInstruction as Content; + // Returns false if the role can't be represented on-device. + if (systemContent.role && systemContent.role === 'function') { + return false; + } + + // Returns false if the system prompt is multi-part. + if (systemContent.parts && systemContent.parts.length > 1) { + return false; + } + + // Returns false if the system prompt isn't text. + const systemText = request.systemInstruction as TextPart; + if (!systemText.text) { + return false; + } + } + + // Returns false if the prompt is empty. + if (request.contents.length === 0) { + return false; + } + + // Applies the same checks as above, but for each content item. + for (const content of request.contents) { + if (content.role === 'function') { + return false; + } + + if (content.parts.length > 1) { + return false; + } + + if (!content.parts[0].text) { + return false; + } + } + + return true; + } + private download(): void { + if (this.downloadPromise) { + return; + } + this.downloadPromise = this.languageModelProvider + ?.create(this.onDeviceParams) + .then((model: LanguageModel) => { + delete this.downloadPromise; + return model; + }); + return; + } + private static toSystemPrompt( + prompt: string | Content | Part | undefined + ): string | undefined { + if (!prompt) { + return undefined; + } + + if (typeof prompt === 'string') { + return prompt; + } + + const systemContent = prompt as Content; + if ( + systemContent.parts && + systemContent.parts[0] && + systemContent.parts[0].text + ) { + return systemContent.parts[0].text; + } + + const systemPart = prompt as Part; + if (systemPart.text) { + return systemPart.text; + } + + return undefined; + } + private static toOnDeviceRole(role: Role): LanguageModelMessageRole { + return role === 'model' ? 'assistant' : 'user'; + } + private static toInitialPrompts( + contents: Content[] + ): LanguageModelMessageShorthand[] { + return contents.map(c => ({ + role: ChromeAdapter.toOnDeviceRole(c.role), + // Assumes contents have been verified to contain only a single TextPart. + content: c.parts[0].text! + })); + } + private async session( + opts: LanguageModelCreateOptions + ): Promise { + const newSession = await this.languageModelProvider!.create(opts); + if (this.oldSession) { + this.oldSession.destroy(); + } + // Holds session reference, so model isn't unloaded from memory. + this.oldSession = newSession; + return newSession; + } } diff --git a/packages/vertexai/src/types/language-model.ts b/packages/vertexai/src/types/language-model.ts index e564ca467b4..f1dbc7eacdc 100644 --- a/packages/vertexai/src/types/language-model.ts +++ b/packages/vertexai/src/types/language-model.ts @@ -32,7 +32,7 @@ export interface LanguageModel extends EventTarget { ): Promise; destroy(): undefined; } -enum Availability { +export enum Availability { 'unavailable', 'downloadable', 'downloading', @@ -67,7 +67,7 @@ interface LanguageModelMessage { role: LanguageModelMessageRole; content: LanguageModelMessageContent[]; } -interface LanguageModelMessageShorthand { +export interface LanguageModelMessageShorthand { role: LanguageModelMessageRole; content: string; } @@ -75,7 +75,7 @@ interface LanguageModelMessageContent { type: LanguageModelMessageType; content: LanguageModelMessageContentValue; } -type LanguageModelMessageRole = 'system' | 'user' | 'assistant'; +export type LanguageModelMessageRole = 'system' | 'user' | 'assistant'; type LanguageModelMessageType = 'text' | 'image' | 'audio'; type LanguageModelMessageContentValue = | ImageBitmapSource From 1186fb7965af736d8d1bbdd6a5f0fef9b80961cf Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Tue, 1 Apr 2025 17:15:47 -0700 Subject: [PATCH 3/6] Integrate with e2e test app --- e2e/sample-apps/modular.js | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/e2e/sample-apps/modular.js b/e2e/sample-apps/modular.js index 9e943e04494..292a11535a3 100644 --- a/e2e/sample-apps/modular.js +++ b/e2e/sample-apps/modular.js @@ -58,7 +58,12 @@ import { onValue, off } from 'firebase/database'; -import { getGenerativeModel, getVertexAI, VertexAI } from 'firebase/vertexai'; +import { + getGenerativeModel, + getVertexAI, + InferenceMode, + VertexAI +} from 'firebase/vertexai'; import { getDataConnect, DataConnect } from 'firebase/data-connect'; /** @@ -332,6 +337,17 @@ function callDataConnect(app) { console.log('[DATACONNECT] initialized'); } +async function callVertex(app) { + console.log('[VERTEX] start'); + const vertex = getVertexAI(app); + const model = getGenerativeModel(vertex, { + mode: InferenceMode.PREFER_ON_DEVICE + }); + const result = await model.generateContent("What is Roko's Basalisk?"); + console.log(result.response.text()); + console.log('[VERTEX] initialized'); +} + /** * Run smoke tests for all products. * Comment out any products you want to ignore. @@ -353,6 +369,7 @@ async function main() { await callVertexAI(app); callDataConnect(app); await authLogout(app); + await callVertex(app); console.log('DONE'); } From 0d078ecbd461c46b0f2e3e7c4a0df369fd7eaa1d Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Thu, 3 Apr 2025 17:47:35 -0700 Subject: [PATCH 4/6] Test model download logic --- .../src/methods/chrome-adapter.test.ts | 77 +++++++++++++++++++ .../vertexai/src/methods/chrome-adapter.ts | 14 ++-- 2 files changed, 84 insertions(+), 7 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 453720c121d..89b140fcebb 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -20,6 +20,7 @@ import sinonChai from 'sinon-chai'; import chaiAsPromised from 'chai-as-promised'; import { ChromeAdapter } from './chrome-adapter'; import { Availability, LanguageModel } from '../types/language-model'; +import { stub } from 'sinon'; use(sinonChai); use(chaiAsPromised); @@ -56,5 +57,81 @@ describe('ChromeAdapter', () => { }) ).to.be.true; }); + it('returns false and triggers download when model is available after download', async () => { + const languageModelProvider = { + availability: () => Promise.resolve(Availability.downloadable), + create: () => Promise.resolve({}) + } as LanguageModel; + const createStub = stub(languageModelProvider, 'create').resolves( + {} as LanguageModel + ); + const adapter = new ChromeAdapter( + languageModelProvider, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }) + ).to.be.false; + expect(createStub).to.have.been.calledOnce; + }); + it('avoids redundant downloads', async () => { + const languageModelProvider = { + availability: () => Promise.resolve(Availability.downloadable), + create: () => Promise.resolve({}) + } as LanguageModel; + const downloadPromise = new Promise(() => { + /* never resolves */ + }); + const createStub = stub(languageModelProvider, 'create').returns( + downloadPromise + ); + const adapter = new ChromeAdapter(languageModelProvider); + await adapter.isAvailable({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }); + await adapter.isAvailable({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }); + expect(createStub).to.have.been.calledOnce; + }); + it('clears state when download completes', async () => { + const languageModelProvider = { + availability: () => Promise.resolve(Availability.downloadable), + create: () => Promise.resolve({}) + } as LanguageModel; + let resolveDownload; + const downloadPromise = new Promise(resolveCallback => { + resolveDownload = resolveCallback; + }); + const createStub = stub(languageModelProvider, 'create').returns( + downloadPromise + ); + const adapter = new ChromeAdapter(languageModelProvider); + await adapter.isAvailable({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }); + resolveDownload!(); + await adapter.isAvailable({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }); + expect(createStub).to.have.been.calledTwice; + }); + it('returns false when model is never available', async () => { + const languageModelProvider = { + availability: () => Promise.resolve(Availability.unavailable), + create: () => Promise.resolve({}) + } as LanguageModel; + const adapter = new ChromeAdapter( + languageModelProvider, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + }) + ).to.be.false; + }); }); }); diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index e5364971479..8c867f09682 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -37,8 +37,9 @@ import { isChrome } from '@firebase/util'; * and encapsulates logic for detecting when on-device is possible. */ export class ChromeAdapter { - downloadPromise: Promise | undefined; - oldSession: LanguageModel | undefined; + private isDownloading = false; + private downloadPromise: Promise | undefined; + private oldSession: LanguageModel | undefined; constructor( private languageModelProvider?: LanguageModel, private mode?: InferenceMode, @@ -152,16 +153,15 @@ export class ChromeAdapter { return true; } private download(): void { - if (this.downloadPromise) { + if (this.isDownloading) { return; } + this.isDownloading = true; this.downloadPromise = this.languageModelProvider ?.create(this.onDeviceParams) - .then((model: LanguageModel) => { - delete this.downloadPromise; - return model; + .then(() => { + this.isDownloading = false; }); - return; } private static toSystemPrompt( prompt: string | Content | Part | undefined From f62a98f9c365b213f64d2129d174b4f4d86b0995 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Thu, 3 Apr 2025 18:16:32 -0700 Subject: [PATCH 5/6] Test request-based availability checks --- .../src/methods/chrome-adapter.test.ts | 146 ++++++++++++++++-- .../vertexai/src/methods/chrome-adapter.ts | 43 +++--- 2 files changed, 157 insertions(+), 32 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 89b140fcebb..1a72fad354c 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -21,29 +21,155 @@ import chaiAsPromised from 'chai-as-promised'; import { ChromeAdapter } from './chrome-adapter'; import { Availability, LanguageModel } from '../types/language-model'; import { stub } from 'sinon'; +import * as util from '@firebase/util'; use(sinonChai); use(chaiAsPromised); describe('ChromeAdapter', () => { - describe('isOnDeviceRequest', () => { - it('returns true for simple text part', async () => { + describe('isAvailable', () => { + it('returns false if mode is only cloud', async () => { + const adapter = new ChromeAdapter(undefined, 'only_in_cloud'); expect( - ChromeAdapter._isOnDeviceRequest({ - contents: [{ role: 'user', parts: [{ text: 'hi' }] }] + await adapter.isAvailable({ + contents: [] }) - ).to.be.true; + ).to.be.false; }); - it('returns false if contents empty', async () => { + it('returns false if browser is not Chrome', async () => { + const chromeStub = stub(util, 'isChrome').returns(false); + const adapter = new ChromeAdapter(undefined, 'prefer_on_device'); expect( - ChromeAdapter._isOnDeviceRequest({ + await adapter.isAvailable({ contents: [] }) ).to.be.false; + chromeStub.restore(); }); - }); - describe('isAvailable', () => { - it('returns true if a model is available', async () => { + it('returns false if AI API is undefined', async () => { + const adapter = new ChromeAdapter(undefined, 'prefer_on_device'); + expect( + await adapter.isAvailable({ + contents: [] + }) + ).to.be.false; + }); + it('returns false if LanguageModel API is undefined', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [] + }) + ).to.be.false; + }); + it('returns false if request contents empty', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [] + }) + ).to.be.false; + }); + it('returns false if request content has function role', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [ + { + role: 'function', + parts: [] + } + ] + }) + ).to.be.false; + }); + it('returns false if request content has multiple parts', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [ + { + role: 'user', + parts: [{ text: 'a' }, { text: 'b' }] + } + ] + }) + ).to.be.false; + }); + it('returns false if request content has non-text part', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [ + { + role: 'user', + parts: [{ inlineData: { mimeType: 'a', data: 'b' } }] + } + ] + }) + ).to.be.false; + }); + it('returns false if request system instruction has function role', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [], + systemInstruction: { + role: 'function', + parts: [] + } + }) + ).to.be.false; + }); + it('returns false if request system instruction has multiple parts', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [], + systemInstruction: { + role: 'function', + parts: [{ text: 'a' }, { text: 'b' }] + } + }) + ).to.be.false; + }); + it('returns false if request system instruction has non-text part', async () => { + const adapter = new ChromeAdapter( + {} as LanguageModel, + 'prefer_on_device' + ); + expect( + await adapter.isAvailable({ + contents: [], + systemInstruction: { + role: 'function', + parts: [{ inlineData: { mimeType: 'a', data: 'b' } }] + } + }) + ).to.be.false; + }); + it('returns true if model is readily available', async () => { const languageModelProvider = { availability: () => Promise.resolve(Availability.available) } as LanguageModel; diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index 8c867f09682..ff782cc5d88 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -67,7 +67,7 @@ export class ChromeAdapter { return false; } // Returns false if the request can't be run on-device. - if (!ChromeAdapter._isOnDeviceRequest(request)) { + if (!ChromeAdapter.isOnDeviceRequest(request)) { return false; } const availability = await this.languageModelProvider.availability(); @@ -109,27 +109,7 @@ export class ChromeAdapter { }) } as Response; } - // Visible for testing - static _isOnDeviceRequest(request: GenerateContentRequest): boolean { - if (request.systemInstruction) { - const systemContent = request.systemInstruction as Content; - // Returns false if the role can't be represented on-device. - if (systemContent.role && systemContent.role === 'function') { - return false; - } - - // Returns false if the system prompt is multi-part. - if (systemContent.parts && systemContent.parts.length > 1) { - return false; - } - - // Returns false if the system prompt isn't text. - const systemText = request.systemInstruction as TextPart; - if (!systemText.text) { - return false; - } - } - + private static isOnDeviceRequest(request: GenerateContentRequest): boolean { // Returns false if the prompt is empty. if (request.contents.length === 0) { return false; @@ -150,6 +130,25 @@ export class ChromeAdapter { } } + if (request.systemInstruction) { + const systemContent = request.systemInstruction as Content; + // Returns false if the role can't be represented on-device. + if (systemContent.role && systemContent.role === 'function') { + return false; + } + + // Returns false if the system prompt is multi-part. + if (systemContent.parts && systemContent.parts.length > 1) { + return false; + } + + // Returns false if the system prompt isn't text. + const systemText = request.systemInstruction as TextPart; + if (!systemText.text) { + return false; + } + } + return true; } private download(): void { From ac117e5ed53791744ce2c30da200714264df01e8 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Fri, 4 Apr 2025 16:02:21 -0700 Subject: [PATCH 6/6] Test content generation --- .../src/methods/chrome-adapter.test.ts | 64 +++++++++++++++++- .../vertexai/src/methods/chrome-adapter.ts | 66 +++++++------------ packages/vertexai/src/types/language-model.ts | 8 +-- 3 files changed, 89 insertions(+), 49 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 1a72fad354c..8906bf398bf 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -19,9 +19,14 @@ import { expect, use } from 'chai'; import sinonChai from 'sinon-chai'; import chaiAsPromised from 'chai-as-promised'; import { ChromeAdapter } from './chrome-adapter'; -import { Availability, LanguageModel } from '../types/language-model'; +import { + Availability, + LanguageModel, + LanguageModelCreateOptions +} from '../types/language-model'; import { stub } from 'sinon'; import * as util from '@firebase/util'; +import { GenerateContentRequest } from '../types'; use(sinonChai); use(chaiAsPromised); @@ -191,16 +196,18 @@ describe('ChromeAdapter', () => { const createStub = stub(languageModelProvider, 'create').resolves( {} as LanguageModel ); + const onDeviceParams = {} as LanguageModelCreateOptions; const adapter = new ChromeAdapter( languageModelProvider, - 'prefer_on_device' + 'prefer_on_device', + onDeviceParams ); expect( await adapter.isAvailable({ contents: [{ role: 'user', parts: [{ text: 'hi' }] }] }) ).to.be.false; - expect(createStub).to.have.been.calledOnce; + expect(createStub).to.have.been.calledOnceWith(onDeviceParams); }); it('avoids redundant downloads', async () => { const languageModelProvider = { @@ -260,4 +267,55 @@ describe('ChromeAdapter', () => { ).to.be.false; }); }); + describe('generateContentOnDevice', () => { + it('generates content', async () => { + const languageModelProvider = { + create: () => Promise.resolve({}) + } as LanguageModel; + const languageModel = { + prompt: i => Promise.resolve(i) + } as LanguageModel; + const createStub = stub(languageModelProvider, 'create').resolves( + languageModel + ); + const promptOutput = 'hi'; + const promptStub = stub(languageModel, 'prompt').resolves(promptOutput); + const onDeviceParams = { + systemPrompt: 'be yourself' + } as LanguageModelCreateOptions; + const adapter = new ChromeAdapter( + languageModelProvider, + 'prefer_on_device', + onDeviceParams + ); + const request = { + contents: [{ role: 'user', parts: [{ text: 'anything' }] }] + } as GenerateContentRequest; + const response = await adapter.generateContentOnDevice(request); + // Asserts initialization params are proxied. + expect(createStub).to.have.been.calledOnceWith(onDeviceParams); + // Asserts Vertex input type is mapped to Chrome type. + expect(promptStub).to.have.been.calledOnceWith([ + { + role: request.contents[0].role, + content: [ + { + type: 'text', + content: request.contents[0].parts[0].text + } + ] + } + ]); + // Asserts expected output. + expect(await response.json()).to.deep.equal({ + candidates: [ + { + content: { + parts: [{ text: promptOutput }] + } + } + ] + }); + }); + }); }); diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index ff782cc5d88..59f080468c2 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -27,8 +27,9 @@ import { Availability, LanguageModel, LanguageModelCreateOptions, + LanguageModelMessage, LanguageModelMessageRole, - LanguageModelMessageShorthand + LanguageModelMessageContent } from '../types/language-model'; import { isChrome } from '@firebase/util'; @@ -85,17 +86,13 @@ export class ChromeAdapter { async generateContentOnDevice( request: GenerateContentRequest ): Promise { - const initialPrompts = ChromeAdapter.toInitialPrompts(request.contents); - // Assumes validation asserted there is at least one initial prompt. - const prompt = initialPrompts.pop()!; - const systemPrompt = ChromeAdapter.toSystemPrompt( - request.systemInstruction + const session = await this.session( + // TODO: normalize on-device params during construction. + this.onDeviceParams || {} ); - const session = await this.session({ - initialPrompts, - systemPrompt - }); - const text = await session.prompt(prompt.content); + const messages = ChromeAdapter.toLanguageModelMessages(request.contents); + const text = await session.prompt(messages); + console.log(text); return { json: () => Promise.resolve({ @@ -162,45 +159,30 @@ export class ChromeAdapter { this.isDownloading = false; }); } - private static toSystemPrompt( - prompt: string | Content | Part | undefined - ): string | undefined { - if (!prompt) { - return undefined; - } - - if (typeof prompt === 'string') { - return prompt; - } - - const systemContent = prompt as Content; - if ( - systemContent.parts && - systemContent.parts[0] && - systemContent.parts[0].text - ) { - return systemContent.parts[0].text; - } - - const systemPart = prompt as Part; - if (systemPart.text) { - return systemPart.text; - } - - return undefined; - } private static toOnDeviceRole(role: Role): LanguageModelMessageRole { return role === 'model' ? 'assistant' : 'user'; } - private static toInitialPrompts( + private static toLanguageModelMessages( contents: Content[] - ): LanguageModelMessageShorthand[] { + ): LanguageModelMessage[] { return contents.map(c => ({ role: ChromeAdapter.toOnDeviceRole(c.role), - // Assumes contents have been verified to contain only a single TextPart. - content: c.parts[0].text! + content: c.parts.map(ChromeAdapter.toLanguageModelMessageContent) })); } + private static toLanguageModelMessageContent( + part: Part + ): LanguageModelMessageContent { + if (part.text) { + return { + type: 'text', + content: part.text + }; + } + // Assumes contents have been verified to contain only a single TextPart. + // TODO: support other input types + throw new Error('Not yet implemented'); + } private async session( opts: LanguageModelCreateOptions ): Promise { diff --git a/packages/vertexai/src/types/language-model.ts b/packages/vertexai/src/types/language-model.ts index f1dbc7eacdc..88354d0aeec 100644 --- a/packages/vertexai/src/types/language-model.ts +++ b/packages/vertexai/src/types/language-model.ts @@ -56,22 +56,22 @@ interface LanguageModelExpectedInput { type: LanguageModelMessageType; languages?: string[]; } -type LanguageModelPrompt = +export type LanguageModelPrompt = | LanguageModelMessage[] | LanguageModelMessageShorthand[] | string; type LanguageModelInitialPrompts = | LanguageModelMessage[] | LanguageModelMessageShorthand[]; -interface LanguageModelMessage { +export interface LanguageModelMessage { role: LanguageModelMessageRole; content: LanguageModelMessageContent[]; } -export interface LanguageModelMessageShorthand { +interface LanguageModelMessageShorthand { role: LanguageModelMessageRole; content: string; } -interface LanguageModelMessageContent { +export interface LanguageModelMessageContent { type: LanguageModelMessageType; content: LanguageModelMessageContentValue; }