From 25c63978efc3e08d697ae5b2baf37d4be4b111a3 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Tue, 1 Apr 2025 17:14:19 -0700 Subject: [PATCH 1/9] Implement ChromeAdapter class --- packages/util/src/environment.ts | 6 + .../src/methods/chrome-adapter.test.ts | 65 ++++++ .../vertexai/src/methods/chrome-adapter.ts | 185 +++++++++++++++++- 3 files changed, 251 insertions(+), 5 deletions(-) create mode 100644 packages/vertexai/src/methods/chrome-adapter.test.ts 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..77d0402fbec --- /dev/null +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -0,0 +1,65 @@ +/** + * @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 { InferenceMode } from '../types'; + +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 aiProvider = { + languageModel: { + capabilities: () => + Promise.resolve({ + available: 'readily' + } as AILanguageModelCapabilities) + } + } as AI; + const adapter = new ChromeAdapter( + aiProvider, + InferenceMode.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 ca97840cf27..3d1a5b9dd76 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -1,7 +1,29 @@ +/** + * @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 { isChrome } from '@firebase/util'; import { + Content, EnhancedGenerateContentResponse, GenerateContentRequest, - InferenceMode + InferenceMode, + Part, + Role, + TextPart } from '../types'; import { AI, AILanguageModelCreateOptionsWithSystemPrompt } from '../types/ai'; @@ -10,22 +32,175 @@ import { AI, AILanguageModelCreateOptionsWithSystemPrompt } from '../types/ai'; * and encapsulates logic for detecting when on-device is possible. */ export class ChromeAdapter { + downloadPromise: Promise | undefined; + oldSession: AILanguageModel | undefined; constructor( private aiProvider?: AI, private mode?: InferenceMode, private onDeviceParams?: AILanguageModelCreateOptionsWithSystemPrompt ) {} - // 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 === InferenceMode.ONLY_ON_CLOUD) { + return false; + } + // Returns false if the on-device inference API is undefined. + const isLanguageModelAvailable = + isChrome() && this.aiProvider && this.aiProvider.languageModel; + if (!isLanguageModelAvailable) { + return false; + } + // Returns false if the request can't be run on-device. + if (!ChromeAdapter._isOnDeviceRequest(request)) { + return false; + } + switch (await this.availability()) { + case 'readily': + // Returns true only if a model is immediately available. + return true; + case 'after-download': + // Triggers async model download. + this.download(); + case 'no': + 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 result = await session.prompt(prompt.content); return { - text: () => '', + text: () => result, functionCalls: () => undefined }; } + // 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 async availability(): Promise { + return this.aiProvider?.languageModel + .capabilities() + .then((c: AILanguageModelCapabilities) => c.available); + } + private download(): void { + if (this.downloadPromise) { + return; + } + this.downloadPromise = this.aiProvider?.languageModel + .create(this.onDeviceParams) + .then((model: AILanguageModel) => { + 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): AILanguageModelPromptRole { + return role === 'model' ? 'assistant' : 'user'; + } + private static toInitialPrompts( + contents: Content[] + ): AILanguageModelPrompt[] { + 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: AILanguageModelCreateOptionsWithSystemPrompt + ): Promise { + const newSession = await this.aiProvider!.languageModel.create(opts); + if (this.oldSession) { + this.oldSession.destroy(); + } + // Holds session reference, so model isn't unloaded from memory. + this.oldSession = newSession; + return newSession; + } } From 15d114c71ce3bbd64476d1118ac314369729fa45 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Tue, 1 Apr 2025 17:15:47 -0700 Subject: [PATCH 2/9] Integrate with e2e test app --- e2e/sample-apps/modular.js | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/e2e/sample-apps/modular.js b/e2e/sample-apps/modular.js index 9e943e04494..2e6544eab6a 100644 --- a/e2e/sample-apps/modular.js +++ b/e2e/sample-apps/modular.js @@ -58,7 +58,7 @@ 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 +332,15 @@ 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 +362,7 @@ async function main() { await callVertexAI(app); callDataConnect(app); await authLogout(app); + await callVertex(app); console.log('DONE'); } From 3f02db006a07a42337d1dae671b50cfe0d4c2a5c Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Thu, 3 Apr 2025 14:01:12 -0700 Subject: [PATCH 3/9] Run yarn format --- e2e/sample-apps/modular.js | 17 ++++++++++++----- packages/vertexai/src/api.test.ts | 6 +----- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/e2e/sample-apps/modular.js b/e2e/sample-apps/modular.js index 2e6544eab6a..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, InferenceMode, VertexAI } from 'firebase/vertexai'; +import { + getGenerativeModel, + getVertexAI, + InferenceMode, + VertexAI +} from 'firebase/vertexai'; import { getDataConnect, DataConnect } from 'firebase/data-connect'; /** @@ -334,10 +339,12 @@ function callDataConnect(app) { 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()) + 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'); } diff --git a/packages/vertexai/src/api.test.ts b/packages/vertexai/src/api.test.ts index 0f25384071a..2852e9c3f1f 100644 --- a/packages/vertexai/src/api.test.ts +++ b/packages/vertexai/src/api.test.ts @@ -14,11 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -import { - ImagenModelParams, - ModelParams, - VertexAIErrorCode -} from './types'; +import { ImagenModelParams, ModelParams, VertexAIErrorCode } from './types'; import { VertexAIError } from './errors'; import { ImagenModel, getGenerativeModel, getImagenModel } from './api'; import { expect } from 'chai'; From ee4d3aabfb7fb9b442a2f988e7ff250556c55ecf Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Thu, 3 Apr 2025 17:47:35 -0700 Subject: [PATCH 4/9] Test model download logic --- .../src/methods/chrome-adapter.test.ts | 94 ++++++++++++++++++- .../vertexai/src/methods/chrome-adapter.ts | 15 +-- 2 files changed, 100 insertions(+), 9 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 45f8d661d7f..a6267faa967 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -19,6 +19,7 @@ import { expect, use } from 'chai'; import sinonChai from 'sinon-chai'; import chaiAsPromised from 'chai-as-promised'; import { ChromeAdapter } from './chrome-adapter'; +import { stub } from 'sinon'; use(sinonChai); use(chaiAsPromised); @@ -41,13 +42,13 @@ describe('ChromeAdapter', () => { }); }); describe('isAvailable', () => { - it('returns true if a model is available', async () => { + it('returns true if model is readily available', async () => { const aiProvider = { languageModel: { capabilities: () => Promise.resolve({ available: 'readily' - } as AILanguageModelCapabilities) + }) } } as AI; const adapter = new ChromeAdapter(aiProvider, 'prefer_on_device'); @@ -57,5 +58,94 @@ describe('ChromeAdapter', () => { }) ).to.be.true; }); + it('returns false and triggers download when model is available after download', async () => { + const aiProvider = { + languageModel: { + capabilities: () => + Promise.resolve({ + available: 'after-download' + }), + create: () => Promise.resolve({}) + } + } as AI; + const createStub = stub(aiProvider.languageModel, 'create').resolves( + {} as AILanguageModel + ); + const adapter = new ChromeAdapter(aiProvider, '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 aiProvider = { + languageModel: { + capabilities: () => + Promise.resolve({ + available: 'after-download' + }), + create: () => {} + } + } as AI; + const downloadPromise = new Promise(() => { + /* never resolves */ + }); + const createStub = stub(aiProvider.languageModel, 'create').returns( + downloadPromise + ); + const adapter = new ChromeAdapter(aiProvider); + 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 aiProvider = { + languageModel: { + capabilities: () => + Promise.resolve({ + available: 'after-download' + }), + create: () => {} + } + } as AI; + let resolveDownload; + const downloadPromise = new Promise(resolveCallback => { + resolveDownload = resolveCallback; + }); + const createStub = stub(aiProvider.languageModel, 'create').returns( + downloadPromise + ); + const adapter = new ChromeAdapter(aiProvider); + 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 aiProvider = { + languageModel: { + capabilities: () => + Promise.resolve({ + available: 'no' + }) + } + } as AI; + const adapter = new ChromeAdapter(aiProvider, '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 a8741068879..fb205cab393 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -31,10 +31,12 @@ import { * and encapsulates logic for detecting when on-device is possible. */ export class ChromeAdapter { - downloadPromise: Promise | undefined; - oldSession: AILanguageModel | undefined; + private isDownloading = false; + private downloadPromise: Promise | undefined; + private oldSession: AILanguageModel | undefined; constructor( private aiProvider?: AI, + // TODO: mode can be required now. private mode?: InferenceMode, private onDeviceParams?: AILanguageModelCreateOptionsWithSystemPrompt ) {} @@ -141,16 +143,15 @@ export class ChromeAdapter { .then((c: AILanguageModelCapabilities) => c.available); } private download(): void { - if (this.downloadPromise) { + if (this.isDownloading) { return; } + this.isDownloading = true; this.downloadPromise = this.aiProvider?.languageModel .create(this.onDeviceParams) - .then((model: AILanguageModel) => { - delete this.downloadPromise; - return model; + .then(() => { + this.isDownloading = false; }); - return; } private static toSystemPrompt( prompt: string | Content | Part | undefined From d546aec1e69c642a4af701fd6347536e6ff6157f Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Thu, 3 Apr 2025 18:16:32 -0700 Subject: [PATCH 5/9] Test request-based availability checks --- .../src/methods/chrome-adapter.test.ts | 120 ++++++++++++++++-- .../vertexai/src/methods/chrome-adapter.ts | 43 +++---- 2 files changed, 132 insertions(+), 31 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index a6267faa967..6f5076b53ed 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -20,28 +20,130 @@ import sinonChai from 'sinon-chai'; import chaiAsPromised from 'chai-as-promised'; import { ChromeAdapter } from './chrome-adapter'; 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(); + }); + 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 AI, 'prefer_on_device'); + expect( + await adapter.isAvailable({ + contents: [] + }) + ).to.be.false; + }); + it('returns false if request contents empty', async () => { + const adapter = new ChromeAdapter({} as AI, '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 AI, '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 AI, '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 AI, '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 AI, '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 AI, '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 AI, 'prefer_on_device'); + expect( + await adapter.isAvailable({ + contents: [], + systemInstruction: { + role: 'function', + parts: [{ inlineData: { mimeType: 'a', data: 'b' } }] + } + }) + ).to.be.false; }); - }); - describe('isAvailable', () => { it('returns true if model is readily available', async () => { const aiProvider = { languageModel: { diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index fb205cab393..0912a9163e9 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -60,7 +60,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; } switch (await this.availability()) { @@ -94,27 +94,7 @@ export class ChromeAdapter { functionCalls: () => undefined }; } - // 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; @@ -135,6 +115,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 async availability(): Promise { From 87fbb9bf73c732814023b3d2cee20812d330c668 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Fri, 4 Apr 2025 16:02:21 -0700 Subject: [PATCH 6/9] Test content generation --- .../src/methods/chrome-adapter.test.ts | 72 ++++++++++++++++++- .../vertexai/src/methods/chrome-adapter.ts | 47 +++--------- 2 files changed, 78 insertions(+), 41 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 6f5076b53ed..48633f0e86b 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -173,13 +173,18 @@ describe('ChromeAdapter', () => { const createStub = stub(aiProvider.languageModel, 'create').resolves( {} as AILanguageModel ); - const adapter = new ChromeAdapter(aiProvider, 'prefer_on_device'); + const onDeviceParams = {} as AILanguageModelCreateOptionsWithSystemPrompt; + const adapter = new ChromeAdapter( + aiProvider, + '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 aiProvider = { @@ -250,4 +255,67 @@ describe('ChromeAdapter', () => { ).to.be.false; }); }); + describe('generateContentOnDevice', () => { + it('Extracts and concats initial prompts', async () => { + const aiProvider = { + languageModel: { + create: () => Promise.resolve({}) + } + } as AI; + const factoryStub = stub(aiProvider.languageModel, 'create').resolves({ + prompt: s => Promise.resolve(s) + } as AILanguageModel); + const text = ['first', 'second', 'third']; + const onDeviceParams = { + initialPrompts: [{ role: 'user', content: text[0] }] + } as AILanguageModelCreateOptionsWithSystemPrompt; + const adapter = new ChromeAdapter( + aiProvider, + 'prefer_on_device', + onDeviceParams + ); + const response = await adapter.generateContentOnDevice({ + contents: [ + { role: 'model', parts: [{ text: text[1] }] }, + { role: 'user', parts: [{ text: text[2] }] } + ] + }); + expect(factoryStub).to.have.been.calledOnceWith({ + initialPrompts: [ + { role: 'user', content: text[0] }, + // Asserts tail is passed as initial prompts, and + // role is normalized from model to assistant. + { role: 'assistant', content: text[1] } + ] + }); + expect(response.text()).to.equal(text[2]); + }); + it('Extracts system prompt', async () => { + const aiProvider = { + languageModel: { + create: () => Promise.resolve({}) + } + } as AI; + const factoryStub = stub(aiProvider.languageModel, 'create').resolves({ + prompt: s => Promise.resolve(s) + } as AILanguageModel); + const onDeviceParams = { + systemPrompt: 'be yourself' + } as AILanguageModelCreateOptionsWithSystemPrompt; + const adapter = new ChromeAdapter( + aiProvider, + 'prefer_on_device', + onDeviceParams + ); + const text = 'hi'; + const response = await adapter.generateContentOnDevice({ + contents: [{ role: 'user', parts: [{ text }] }] + }); + expect(factoryStub).to.have.been.calledOnceWith({ + initialPrompts: [], + systemPrompt: onDeviceParams.systemPrompt + }); + expect(response.text()).to.equal(text); + }); + }); }); diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index 0912a9163e9..0d90d470284 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -21,7 +21,6 @@ import { EnhancedGenerateContentResponse, GenerateContentRequest, InferenceMode, - Part, Role, TextPart } from '../types'; @@ -78,16 +77,13 @@ export class ChromeAdapter { async generateContentOnDevice( request: GenerateContentRequest ): Promise { - const initialPrompts = ChromeAdapter.toInitialPrompts(request.contents); + const createOptions = this.onDeviceParams || {}; + createOptions.initialPrompts ??= []; + const extractedInitialPrompts = 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 prompt = extractedInitialPrompts.pop()!; + createOptions.initialPrompts.push(...extractedInitialPrompts); + const session = await this.session(createOptions); const result = await session.prompt(prompt.content); return { text: () => result, @@ -152,33 +148,6 @@ 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): AILanguageModelPromptRole { return role === 'model' ? 'assistant' : 'user'; } @@ -192,9 +161,9 @@ export class ChromeAdapter { })); } private async session( - opts: AILanguageModelCreateOptionsWithSystemPrompt + options: AILanguageModelCreateOptionsWithSystemPrompt ): Promise { - const newSession = await this.aiProvider!.languageModel.create(opts); + const newSession = await this.aiProvider!.languageModel.create(options); if (this.oldSession) { this.oldSession.destroy(); } From 29d399f4b5465457ff558bbe2882607805c96140 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Fri, 4 Apr 2025 16:19:44 -0700 Subject: [PATCH 7/9] Remove request.systemInstruction validation We only define system prompts via onDeviceParams initialization. --- .../src/methods/chrome-adapter.test.ts | 36 ------------------- .../vertexai/src/methods/chrome-adapter.ts | 22 +----------- 2 files changed, 1 insertion(+), 57 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 48633f0e86b..0818988894e 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -108,42 +108,6 @@ describe('ChromeAdapter', () => { }) ).to.be.false; }); - it('returns false if request system instruction has function role', async () => { - const adapter = new ChromeAdapter({} as AI, '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 AI, '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 AI, '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 aiProvider = { languageModel: { diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index 0d90d470284..bb4611d8557 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -21,8 +21,7 @@ import { EnhancedGenerateContentResponse, GenerateContentRequest, InferenceMode, - Role, - TextPart + Role } from '../types'; /** @@ -111,25 +110,6 @@ 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 async availability(): Promise { From 1b7233d032ea8c71f3ce7be33bd2db4677577e20 Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Mon, 7 Apr 2025 17:09:05 -0700 Subject: [PATCH 8/9] Refactor to emulate Vertex response --- .../src/methods/chrome-adapter.test.ts | 20 ++++++++++++++-- .../vertexai/src/methods/chrome-adapter.ts | 23 ++++++++++++++----- .../src/methods/generate-content.test.ts | 15 ++++++------ .../vertexai/src/methods/generate-content.ts | 16 ++++++------- 4 files changed, 49 insertions(+), 25 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.test.ts b/packages/vertexai/src/methods/chrome-adapter.test.ts index 0818988894e..c350bd366c0 100644 --- a/packages/vertexai/src/methods/chrome-adapter.test.ts +++ b/packages/vertexai/src/methods/chrome-adapter.test.ts @@ -252,7 +252,15 @@ describe('ChromeAdapter', () => { { role: 'assistant', content: text[1] } ] }); - expect(response.text()).to.equal(text[2]); + expect(await response.json()).to.deep.equal({ + candidates: [ + { + content: { + parts: [{ text: text[2] }] + } + } + ] + }); }); it('Extracts system prompt', async () => { const aiProvider = { @@ -279,7 +287,15 @@ describe('ChromeAdapter', () => { initialPrompts: [], systemPrompt: onDeviceParams.systemPrompt }); - expect(response.text()).to.equal(text); + expect(await response.json()).to.deep.equal({ + candidates: [ + { + content: { + parts: [{ text }] + } + } + ] + }); }); }); }); diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index bb4611d8557..312660b1da4 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -18,7 +18,6 @@ import { isChrome } from '@firebase/util'; import { Content, - EnhancedGenerateContentResponse, GenerateContentRequest, InferenceMode, Role @@ -75,19 +74,31 @@ export class ChromeAdapter { } async generateContentOnDevice( request: GenerateContentRequest - ): Promise { + ): Promise { const createOptions = this.onDeviceParams || {}; createOptions.initialPrompts ??= []; - const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(request.contents); + const extractedInitialPrompts = ChromeAdapter.toInitialPrompts( + request.contents + ); // Assumes validation asserted there is at least one initial prompt. const prompt = extractedInitialPrompts.pop()!; createOptions.initialPrompts.push(...extractedInitialPrompts); const session = await this.session(createOptions); const result = await session.prompt(prompt.content); + return ChromeAdapter.toResponse(result); + } + private static toResponse(text: string): Response { return { - text: () => result, - functionCalls: () => undefined - }; + json: async () => ({ + candidates: [ + { + content: { + parts: [{ text }] + } + } + ] + }) + } as Response; } private static isOnDeviceRequest(request: GenerateContentRequest): boolean { // Returns false if the prompt is empty. diff --git a/packages/vertexai/src/methods/generate-content.test.ts b/packages/vertexai/src/methods/generate-content.test.ts index ec825421709..3a6cffb9f9d 100644 --- a/packages/vertexai/src/methods/generate-content.test.ts +++ b/packages/vertexai/src/methods/generate-content.test.ts @@ -291,24 +291,23 @@ describe('generateContent()', () => { expect(mockFetch).to.be.called; }); it('on-device', async () => { - const expectedText = 'hi'; const chromeAdapter = new ChromeAdapter(); const mockIsAvailable = stub(chromeAdapter, 'isAvailable').resolves(true); - const mockGenerateContent = stub( + const mockResponse = getMockResponse( + 'unary-success-basic-reply-short.json' + ); + const makeRequestStub = stub( chromeAdapter, 'generateContentOnDevice' - ).resolves({ - text: () => expectedText, - functionCalls: () => undefined - }); + ).resolves(mockResponse as Response); const result = await generateContent( fakeApiSettings, 'model', fakeRequestParams, chromeAdapter ); - expect(result.response.text()).to.equal(expectedText); + expect(result.response.text()).to.include('Mountain View, California'); expect(mockIsAvailable).to.be.called; - expect(mockGenerateContent).to.be.calledWith(fakeRequestParams); + expect(makeRequestStub).to.be.calledWith(fakeRequestParams); }); }); diff --git a/packages/vertexai/src/methods/generate-content.ts b/packages/vertexai/src/methods/generate-content.ts index 63745c47fae..ba7a162aa9c 100644 --- a/packages/vertexai/src/methods/generate-content.ts +++ b/packages/vertexai/src/methods/generate-content.ts @@ -16,7 +16,6 @@ */ import { - EnhancedGenerateContentResponse, GenerateContentRequest, GenerateContentResponse, GenerateContentResult, @@ -51,8 +50,8 @@ async function generateContentOnCloud( model: string, params: GenerateContentRequest, requestOptions?: RequestOptions -): Promise { - const response = await makeRequest( +): Promise { + return makeRequest( model, Task.GENERATE_CONTENT, apiSettings, @@ -60,9 +59,6 @@ async function generateContentOnCloud( JSON.stringify(params), requestOptions ); - const responseJson: GenerateContentResponse = await response.json(); - const enhancedResponse = createEnhancedContentResponse(responseJson); - return enhancedResponse; } export async function generateContent( @@ -72,17 +68,19 @@ export async function generateContent( chromeAdapter: ChromeAdapter, requestOptions?: RequestOptions ): Promise { - let enhancedResponse; + let response; if (await chromeAdapter.isAvailable(params)) { - enhancedResponse = await chromeAdapter.generateContentOnDevice(params); + response = await chromeAdapter.generateContentOnDevice(params); } else { - enhancedResponse = await generateContentOnCloud( + response = await generateContentOnCloud( apiSettings, model, params, requestOptions ); } + const responseJson: GenerateContentResponse = await response.json(); + const enhancedResponse = createEnhancedContentResponse(responseJson); return { response: enhancedResponse }; From 4db0d63b4cc32a4bc7ef41672870df7f5af8f81a Mon Sep 17 00:00:00 2001 From: Erik Eldridge Date: Tue, 8 Apr 2025 08:51:46 -0700 Subject: [PATCH 9/9] Update content generator to emulate Vertex response --- .../vertexai/src/methods/chrome-adapter.ts | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/packages/vertexai/src/methods/chrome-adapter.ts b/packages/vertexai/src/methods/chrome-adapter.ts index f66e02f6711..3c9fc89e88f 100644 --- a/packages/vertexai/src/methods/chrome-adapter.ts +++ b/packages/vertexai/src/methods/chrome-adapter.ts @@ -15,11 +15,7 @@ * limitations under the License. */ -import { - EnhancedGenerateContentResponse, - GenerateContentRequest, - InferenceMode -} from '../types'; +import { GenerateContentRequest, InferenceMode } from '../types'; /** * Defines an inference "backend" that uses Chrome's on-device model, @@ -38,10 +34,18 @@ export class ChromeAdapter { async generateContentOnDevice( // eslint-disable-next-line @typescript-eslint/no-unused-vars request: GenerateContentRequest - ): Promise { + ): Promise { return { - text: () => '', - functionCalls: () => undefined - }; + json: () => + Promise.resolve({ + candidates: [ + { + content: { + parts: [{ text: '' }] + } + } + ] + }) + } as Response; } }