Complete Hybrid inference impl

Fix languageCode parameter in action_code_url (#8912)

* Fix languageCode parameter in action_code_url

* Add changeset

Vaihi add langmodel types. (#8927)

* Adding LanguageModel types. These are based off https://github.com/webmachinelearning/prompt-api?tab=readme-ov-file#full-api-surface-in-web-idl

* Adding LanguageModel types.

* Remove bunch of exports

* yarn formatted

* after lint

Define HybridParams (#8935)

Co-authored-by: Erik Eldridge <[email protected]>

Adding smoke test for new hybrid params (#8937)

* Adding smoke test for new hybrid params

* Use the existing name of the model params input

---------

Co-authored-by: Erik Eldridge <[email protected]>

Moving to in-cloud naming (#8938)

Co-authored-by: Erik Eldridge <[email protected]>

Moving to string type for the inference mode (#8941)

Define ChromeAdapter class (#8942)

Co-authored-by: Erik Eldridge <[email protected]>

VinF Hybrid Inference: Implement ChromeAdapter (rebased) (#8943)

Adding count token impl (#8950)

VinF Hybrid Inference #4: ChromeAdapter in stream methods (rebased) (#8949)

Define values for Availability enum (#8951)

VinF Hybrid Inference: narrow Chrome input type (#8953)

Add image inference support (#8954)

* Adding image based input for inference

* adding image as input to create language model object

disable count tokens api for on-device inference (#8962)

VinF Hybrid Inference: throw if only_on_device and model is unavailable (#8965)

Author: pashanka

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 } from 'firebase/vertexai';
 import { getDataConnect, DataConnect } from 'firebase/data-connect';
 
 /**
@@ -313,9 +313,15 @@ function callPerformance(app) {
 async function callVertexAI(app) {
   console.log('[VERTEXAI] start');
   const vertexAI = getVertexAI(app);
-  const model = getGenerativeModel(vertexAI, { model: 'gemini-1.5-flash' });
-  const result = await model.countTokens('abcdefg');
-  console.log(`[VERTEXAI] counted tokens: ${result.totalTokens}`);
+  const model = getGenerativeModel(vertexAI, {
+    mode: 'only_on_device'
+  });
+  const singleResult = await model.generateContent([
+    { text: 'describe the following:' },
+    { text: 'the mojave desert' }
+  ]);
+  console.log(`Generated text: ${singleResult.response.text()}`);
+  console.log(`[VERTEXAI] end`);
 }
 
 /**
@@ -341,18 +347,18 @@ async function main() {
   const app = initializeApp(config);
   setLogLevel('warn');
 
-  callAppCheck(app);
-  await authLogin(app);
-  await callStorage(app);
-  await callFirestore(app);
-  await callDatabase(app);
-  await callMessaging(app);
-  callAnalytics(app);
-  callPerformance(app);
-  await callFunctions(app);
+  // callAppCheck(app);
+  // await authLogin(app);
+  // await callStorage(app);
+  // await callFirestore(app);
+  // await callDatabase(app);
+  // await callMessaging(app);
+  // callAnalytics(app);
+  // callPerformance(app);
+  // await callFunctions(app);
   await callVertexAI(app);
-  callDataConnect(app);
-  await authLogout(app);
+  // callDataConnect(app);
+  // await authLogout(app);
   console.log('DONE');
 }
 

packages/vertexai/src/api.test.ts

@@ -101,6 +101,21 @@ describe('Top level API', () => {
     expect(genModel).to.be.an.instanceOf(GenerativeModel);
     expect(genModel.model).to.equal('publishers/google/models/my-model');
   });
+  it('getGenerativeModel with HybridParams sets a default model', () => {
+    const genModel = getGenerativeModel(fakeAI, {
+      mode: 'only_on_device'
+    });
+    expect(genModel.model).to.equal(
+      `publishers/google/models/${GenerativeModel.DEFAULT_HYBRID_IN_CLOUD_MODEL}`
+    );
+  });
+  it('getGenerativeModel with HybridParams honors a model override', () => {
+    const genModel = getGenerativeModel(fakeAI, {
+      mode: 'prefer_on_device',
+      inCloudParams: { model: 'my-model' }
+    });
+    expect(genModel.model).to.equal('publishers/google/models/my-model');
+  });
   it('getImagenModel throws if no model is provided', () => {
     try {
       getImagenModel(fakeAI, {} as ImagenModelParams);

packages/vertexai/src/api.ts

@@ -23,6 +23,7 @@ import { AIService } from './service';
 import { AI, AIOptions, VertexAI, VertexAIOptions } from './public-types';
 import {
   ImagenModelParams,
+  HybridParams,
   ModelParams,
   RequestOptions,
   AIErrorCode
@@ -31,6 +32,8 @@ import { AIError } from './errors';
 import { AIModel, GenerativeModel, ImagenModel } from './models';
 import { encodeInstanceIdentifier } from './helpers';
 import { GoogleAIBackend, VertexAIBackend } from './backend';
+import { ChromeAdapter } from './methods/chrome-adapter';
+import { LanguageModel } from './types/language-model';
 
 export { ChatSession } from './methods/chat-session';
 export * from './requests/schema-builder';
@@ -138,16 +141,36 @@ export function getAI(
  */
 export function getGenerativeModel(
   ai: AI,
-  modelParams: ModelParams,
+  modelParams: ModelParams | HybridParams,
   requestOptions?: RequestOptions
 ): GenerativeModel {
-  if (!modelParams.model) {
+  // Uses the existence of HybridParams.mode to clarify the type of the modelParams input.
+  const hybridParams = modelParams as HybridParams;
+  let inCloudParams: ModelParams;
+  if (hybridParams.mode) {
+    inCloudParams = hybridParams.inCloudParams || {
+      model: GenerativeModel.DEFAULT_HYBRID_IN_CLOUD_MODEL
+    };
+  } else {
+    inCloudParams = modelParams as ModelParams;
+  }
+
+  if (!inCloudParams.model) {
     throw new AIError(
       AIErrorCode.NO_MODEL,
       `Must provide a model name. Example: getGenerativeModel({ model: 'my-model-name' })`
     );
   }
-  return new GenerativeModel(ai, modelParams, requestOptions);
+  return new GenerativeModel(
+    ai,
+    inCloudParams,
+    new ChromeAdapter(
+      window.LanguageModel as LanguageModel,
+      hybridParams.mode,
+      hybridParams.onDeviceParams
+    ),
+    requestOptions
+  );
 }
 
 /**

packages/vertexai/src/methods/chat-session.test.ts

@@ -24,6 +24,7 @@ import { GenerateContentStreamResult } from '../types';
 import { ChatSession } from './chat-session';
 import { ApiSettings } from '../types/internal';
 import { VertexAIBackend } from '../backend';
+import { ChromeAdapter } from './chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -46,7 +47,11 @@ describe('ChatSession', () => {
         generateContentMethods,
         'generateContent'
       ).rejects('generateContent failed');
-      const chatSession = new ChatSession(fakeApiSettings, 'a-model');
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        new ChromeAdapter()
+      );
       await expect(chatSession.sendMessage('hello')).to.be.rejected;
       expect(generateContentStub).to.be.calledWith(
         fakeApiSettings,
@@ -63,7 +68,11 @@ describe('ChatSession', () => {
         generateContentMethods,
         'generateContentStream'
       ).rejects('generateContentStream failed');
-      const chatSession = new ChatSession(fakeApiSettings, 'a-model');
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        new ChromeAdapter()
+      );
       await expect(chatSession.sendMessageStream('hello')).to.be.rejected;
       expect(generateContentStreamStub).to.be.calledWith(
         fakeApiSettings,
@@ -82,7 +91,11 @@ describe('ChatSession', () => {
         generateContentMethods,
         'generateContentStream'
       ).resolves({} as unknown as GenerateContentStreamResult);
-      const chatSession = new ChatSession(fakeApiSettings, 'a-model');
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        new ChromeAdapter()
+      );
       await chatSession.sendMessageStream('hello');
       expect(generateContentStreamStub).to.be.calledWith(
         fakeApiSettings,

packages/vertexai/src/methods/chat-session.ts

@@ -30,6 +30,7 @@ import { validateChatHistory } from './chat-session-helpers';
 import { generateContent, generateContentStream } from './generate-content';
 import { ApiSettings } from '../types/internal';
 import { logger } from '../logger';
+import { ChromeAdapter } from './chrome-adapter';
 
 /**
  * Do not log a message for this error.
@@ -50,6 +51,7 @@ export class ChatSession {
   constructor(
     apiSettings: ApiSettings,
     public model: string,
+    private chromeAdapter: ChromeAdapter,
     public params?: StartChatParams,
     public requestOptions?: RequestOptions
   ) {
@@ -95,6 +97,7 @@ export class ChatSession {
           this._apiSettings,
           this.model,
           generateContentRequest,
+          this.chromeAdapter,
           this.requestOptions
         )
       )
@@ -146,6 +149,7 @@ export class ChatSession {
       this._apiSettings,
       this.model,
       generateContentRequest,
+      this.chromeAdapter,
       this.requestOptions
     );
 

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

@@ -0,0 +1,287 @@
+/**
+ * @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 { AIError } from '../errors';
+import {
+  CountTokensRequest,
+  GenerateContentRequest,
+  InferenceMode,
+  Part,
+  AIErrorCode
+} from '../types';
+import {
+  Availability,
+  LanguageModel,
+  LanguageModelCreateOptions,
+  LanguageModelMessageContent
+} from '../types/language-model';
+
+/**
+ * Defines an inference "backend" that uses Chrome's on-device model,
+ * and encapsulates logic for detecting when on-device is possible.
+ */
+export class ChromeAdapter {
+  private isDownloading = false;
+  private downloadPromise: Promise<LanguageModel | void> | undefined;
+  private oldSession: LanguageModel | undefined;
+  constructor(
+    private languageModelProvider?: LanguageModel,
+    private mode?: InferenceMode,
+    private onDeviceParams?: LanguageModelCreateOptions
+  ) {}
+
+  /**
+   * Checks if a given request can be made on-device.
+   *
+   * <ol>Encapsulates a few concerns:
+   *   <li>the mode</li>
+   *   <li>API existence</li>
+   *   <li>prompt formatting</li>
+   *   <li>model availability, including triggering download if necessary</li>
+   * </ol>
+   *
+   * <p>Pros: callers needn't be concerned with details of on-device availability.</p>
+   * <p>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.</p>
+   */
+  async isAvailable(request: GenerateContentRequest): Promise<boolean> {
+    if (this.mode === 'only_in_cloud') {
+      return false;
+    }
+
+    const availability = await this.languageModelProvider?.availability();
+
+    // Triggers async model download so it'll be available next time.
+    if (availability === Availability.downloadable) {
+      this.download();
+    }
+
+    if (this.mode === 'only_on_device') {
+      return true;
+    }
+
+    // Applies prefer_on_device logic.
+    return (
+      availability === Availability.available &&
+      ChromeAdapter.isOnDeviceRequest(request)
+    );
+  }
+
+  /**
+   * Generates content on device.
+   *
+   * <p>This is comparable to {@link GenerativeModel.generateContent} for generating content in
+   * Cloud.</p>
+   * @param request a standard Vertex {@link GenerateContentRequest}
+   * @returns {@link Response}, so we can reuse common response formatting.
+   */
+  async generateContent(request: GenerateContentRequest): Promise<Response> {
+    const session = await this.createSession(
+      // TODO: normalize on-device params during construction.
+      this.onDeviceParams || {}
+    );
+    // TODO: support multiple content objects when Chrome supports
+    // sequence<LanguageModelMessage>
+    const contents = await Promise.all(
+      request.contents[0].parts.map(ChromeAdapter.toLanguageModelMessageContent)
+    );
+    const text = await session.prompt(contents);
+    return ChromeAdapter.toResponse(text);
+  }
+
+  /**
+   * Generates content stream on device.
+   *
+   * <p>This is comparable to {@link GenerativeModel.generateContentStream} for generating content in
+   * Cloud.</p>
+   * @param request a standard Vertex {@link GenerateContentRequest}
+   * @returns {@link Response}, so we can reuse common response formatting.
+   */
+  async generateContentStream(
+    request: GenerateContentRequest
+  ): Promise<Response> {
+    const session = await this.createSession(
+      // TODO: normalize on-device params during construction.
+      this.onDeviceParams || {}
+    );
+    // TODO: support multiple content objects when Chrome supports
+    // sequence<LanguageModelMessage>
+    const contents = await Promise.all(
+      request.contents[0].parts.map(ChromeAdapter.toLanguageModelMessageContent)
+    );
+    const stream = await session.promptStreaming(contents);
+    return ChromeAdapter.toStreamResponse(stream);
+  }
+
+  async countTokens(_request: CountTokensRequest): Promise<Response> {
+    throw new AIError(
+      AIErrorCode.REQUEST_ERROR,
+      'Count Tokens is not yet available for on-device model.'
+    );
+  }
+
+  /**
+   * Asserts inference for the given request can be performed by an on-device model.
+   */
+  private static isOnDeviceRequest(request: GenerateContentRequest): boolean {
+    // 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;
+      }
+    }
+
+    return true;
+  }
+
+  /**
+   * Triggers the download of an on-device model.
+   *
+   * <p>Chrome only downloads models as needed. Chrome knows a model is needed when code calls
+   * LanguageModel.create.</p>
+   *
+   * <p>Since Chrome manages the download, the SDK can only avoid redundant download requests by
+   * tracking if a download has previously been requested.</p>
+   */
+  private download(): void {
+    if (this.isDownloading) {
+      return;
+    }
+    this.isDownloading = true;
+    const options = this.onDeviceParams || {};
+    ChromeAdapter.addImageTypeAsExpectedInput(options);
+    this.downloadPromise = this.languageModelProvider
+      ?.create(options)
+      .then(() => {
+        this.isDownloading = false;
+      });
+  }
+
+  /**
+   * Converts a Vertex Part object to a Chrome LanguageModelMessageContent object.
+   */
+  private static async toLanguageModelMessageContent(
+    part: Part
+  ): Promise<LanguageModelMessageContent> {
+    if (part.text) {
+      return {
+        type: 'text',
+        content: part.text
+      };
+    } else if (part.inlineData) {
+      const formattedImageContent = await fetch(
+        `data:${part.inlineData.mimeType};base64,${part.inlineData.data}`
+      );
+      const imageBlob = await formattedImageContent.blob();
+      const imageBitmap = await createImageBitmap(imageBlob);
+      return {
+        type: 'image',
+        content: imageBitmap
+      };
+    }
+    // Assumes contents have been verified to contain only a single TextPart.
+    // TODO: support other input types
+    throw new Error('Not yet implemented');
+  }
+
+  /**
+   * Abstracts Chrome session creation.
+   *
+   * <p>Chrome uses a multi-turn session for all inference. Vertex uses single-turn for all
+   * inference. To map the Vertex API to Chrome's API, the SDK creates a new session for all
+   * inference.</p>
+   *
+   * <p>Chrome will remove a model from memory if it's no longer in use, so this method ensures a
+   * new session is created before an old session is destroyed.</p>
+   */
+  private async createSession(
+    // TODO: define a default value, since these are optional.
+    options: LanguageModelCreateOptions
+  ): Promise<LanguageModel> {
+    if (!this.languageModelProvider) {
+      throw new AIError(
+        AIErrorCode.REQUEST_ERROR,
+        'Chrome AI requested for unsupported browser version.'
+      );
+    }
+    // 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();
+    }
+    // Holds session reference, so model isn't unloaded from memory.
+    this.oldSession = newSession;
+    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.
+   */
+  private static toResponse(text: string): Response {
+    return {
+      json: async () => ({
+        candidates: [
+          {
+            content: {
+              parts: [{ text }]
+            }
+          }
+        ]
+      })
+    } as Response;
+  }
+
+  /**
+   * Formats string stream returned by Chrome as SSE returned by Vertex.
+   */
+  private static toStreamResponse(stream: ReadableStream<string>): Response {
+    const encoder = new TextEncoder();
+    return {
+      body: stream.pipeThrough(
+        new TransformStream({
+          transform(chunk, controller) {
+            const json = JSON.stringify({
+              candidates: [
+                {
+                  content: {
+                    role: 'model',
+                    parts: [{ text: chunk }]
+                  }
+                }
+              ]
+            });
+            controller.enqueue(encoder.encode(`data: ${json}\n\n`));
+          }
+        })
+      )
+    } as Response;
+  }
+}

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

@@ -27,6 +27,7 @@ import { ApiSettings } from '../types/internal';
 import { Task } from '../requests/request';
 import { mapCountTokensRequest } from '../googleai-mappers';
 import { GoogleAIBackend, VertexAIBackend } from '../backend';
+import { ChromeAdapter } from './chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -66,7 +67,8 @@ describe('countTokens()', () => {
     const result = await countTokens(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.totalTokens).to.equal(6);
     expect(result.totalBillableCharacters).to.equal(16);
@@ -92,7 +94,8 @@ describe('countTokens()', () => {
     const result = await countTokens(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.totalTokens).to.equal(1837);
     expect(result.totalBillableCharacters).to.equal(117);
@@ -120,7 +123,8 @@ describe('countTokens()', () => {
     const result = await countTokens(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.totalTokens).to.equal(258);
     expect(result).to.not.have.property('totalBillableCharacters');
@@ -146,7 +150,12 @@ describe('countTokens()', () => {
       json: mockResponse.json
     } as Response);
     await expect(
-      countTokens(fakeApiSettings, 'model', fakeRequestParams)
+      countTokens(
+        fakeApiSettings,
+        'model',
+        fakeRequestParams,
+        new ChromeAdapter()
+      )
     ).to.be.rejectedWith(/404.*not found/);
     expect(mockFetch).to.be.called;
   });
@@ -164,7 +173,12 @@ describe('countTokens()', () => {
     it('maps request to GoogleAI format', async () => {
       makeRequestStub.resolves({ ok: true, json: () => {} } as Response); // Unused
 
-      await countTokens(fakeGoogleAIApiSettings, 'model', fakeRequestParams);
+      await countTokens(
+        fakeGoogleAIApiSettings,
+        'model',
+        fakeRequestParams,
+        new ChromeAdapter()
+      );
 
       expect(makeRequestStub).to.be.calledWith(
         'model',
@@ -176,4 +190,24 @@ describe('countTokens()', () => {
       );
     });
   });
+  it('on-device', async () => {
+    const chromeAdapter = new ChromeAdapter();
+    const isAvailableStub = stub(chromeAdapter, 'isAvailable').resolves(true);
+    const mockResponse = getMockResponse(
+      'vertexAI',
+      'unary-success-total-tokens.json'
+    );
+    const countTokensStub = stub(chromeAdapter, 'countTokens').resolves(
+      mockResponse as Response
+    );
+    const result = await countTokens(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams,
+      chromeAdapter
+    );
+    expect(result.totalTokens).eq(6);
+    expect(isAvailableStub).to.be.called;
+    expect(countTokensStub).to.be.calledWith(fakeRequestParams);
+  });
 });

packages/vertexai/src/methods/count-tokens.test.ts

@@ -24,8 +24,9 @@ import { Task, makeRequest } from '../requests/request';
 import { ApiSettings } from '../types/internal';
 import * as GoogleAIMapper from '../googleai-mappers';
 import { BackendType } from '../public-types';
+import { ChromeAdapter } from './chrome-adapter';
 
-export async function countTokens(
+export async function countTokensOnCloud(
   apiSettings: ApiSettings,
   model: string,
   params: CountTokensRequest,
@@ -48,3 +49,17 @@ export async function countTokens(
   );
   return response.json();
 }
+
+export async function countTokens(
+  apiSettings: ApiSettings,
+  model: string,
+  params: CountTokensRequest,
+  chromeAdapter: ChromeAdapter,
+  requestOptions?: RequestOptions
+): Promise<CountTokensResponse> {
+  if (await chromeAdapter.isAvailable(params)) {
+    return (await chromeAdapter.countTokens(params)).json();
+  }
+
+  return countTokensOnCloud(apiSettings, model, params, requestOptions);
+}

packages/vertexai/src/methods/count-tokens.ts

@@ -34,6 +34,7 @@ import { Task } from '../requests/request';
 import { AIError } from '../api';
 import { mapGenerateContentRequest } from '../googleai-mappers';
 import { GoogleAIBackend, VertexAIBackend } from '../backend';
+import { ChromeAdapter } from './chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -96,7 +97,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text()).to.include('Mountain View, California');
     expect(makeRequestStub).to.be.calledWith(
@@ -119,7 +121,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text()).to.include('Use Freshly Ground Coffee');
     expect(result.response.text()).to.include('30 minutes of brewing');
@@ -142,7 +145,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.usageMetadata?.totalTokenCount).to.equal(1913);
     expect(result.response.usageMetadata?.candidatesTokenCount).to.equal(76);
@@ -177,7 +181,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text()).to.include(
       'Some information cited from an external source'
@@ -204,7 +209,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text).to.throw('SAFETY');
     expect(makeRequestStub).to.be.calledWith(
@@ -226,7 +232,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text).to.throw('SAFETY');
     expect(makeRequestStub).to.be.calledWith(
@@ -248,7 +255,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text()).to.equal('');
     expect(makeRequestStub).to.be.calledWith(
@@ -270,7 +278,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text()).to.include('Some text');
     expect(makeRequestStub).to.be.calledWith(
@@ -292,7 +301,12 @@ describe('generateContent()', () => {
       json: mockResponse.json
     } as Response);
     await expect(
-      generateContent(fakeApiSettings, 'model', fakeRequestParams)
+      generateContent(
+        fakeApiSettings,
+        'model',
+        fakeRequestParams,
+        new ChromeAdapter()
+      )
     ).to.be.rejectedWith(/400.*invalid argument/);
     expect(mockFetch).to.be.called;
   });
@@ -307,7 +321,12 @@ describe('generateContent()', () => {
       json: mockResponse.json
     } as Response);
     await expect(
-      generateContent(fakeApiSettings, 'model', fakeRequestParams)
+      generateContent(
+        fakeApiSettings,
+        'model',
+        fakeRequestParams,
+        new ChromeAdapter()
+      )
     ).to.be.rejectedWith(
       /firebasevertexai\.googleapis[\s\S]*my-project[\s\S]*api-not-enabled/
     );
@@ -347,7 +366,8 @@ describe('generateContent()', () => {
         generateContent(
           fakeGoogleAIApiSettings,
           'model',
-          requestParamsWithMethod
+          requestParamsWithMethod,
+          new ChromeAdapter()
         )
       ).to.be.rejectedWith(AIError, AIErrorCode.UNSUPPORTED);
       expect(makeRequestStub).to.not.be.called;
@@ -362,7 +382,8 @@ describe('generateContent()', () => {
       await generateContent(
         fakeGoogleAIApiSettings,
         'model',
-        fakeGoogleAIRequestParams
+        fakeGoogleAIRequestParams,
+        new ChromeAdapter()
       );
 
       expect(makeRequestStub).to.be.calledWith(
@@ -375,4 +396,25 @@ describe('generateContent()', () => {
       );
     });
   });
+  // TODO: define a similar test for generateContentStream
+  it('on-device', async () => {
+    const chromeAdapter = new ChromeAdapter();
+    const isAvailableStub = stub(chromeAdapter, 'isAvailable').resolves(true);
+    const mockResponse = getMockResponse(
+      'vertexAI',
+      'unary-success-basic-reply-short.json'
+    );
+    const generateContentStub = stub(chromeAdapter, 'generateContent').resolves(
+      mockResponse as Response
+    );
+    const result = await generateContent(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams,
+      chromeAdapter
+    );
+    expect(result.response.text()).to.include('Mountain View, California');
+    expect(isAvailableStub).to.be.called;
+    expect(generateContentStub).to.be.calledWith(fakeRequestParams);
+  });
 });

packages/vertexai/src/methods/generate-content.test.ts

@@ -28,44 +28,85 @@ import { processStream } from '../requests/stream-reader';
 import { ApiSettings } from '../types/internal';
 import * as GoogleAIMapper from '../googleai-mappers';
 import { BackendType } from '../public-types';
+import { ChromeAdapter } from './chrome-adapter';
 
-export async function generateContentStream(
+async function generateContentStreamOnCloud(
   apiSettings: ApiSettings,
   model: string,
   params: GenerateContentRequest,
   requestOptions?: RequestOptions
-): Promise<GenerateContentStreamResult> {
+): Promise<Response> {
   if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
     params = GoogleAIMapper.mapGenerateContentRequest(params);
   }
-  const response = await makeRequest(
+  return makeRequest(
     model,
     Task.STREAM_GENERATE_CONTENT,
     apiSettings,
     /* stream */ true,
     JSON.stringify(params),
     requestOptions
   );
+}
+
+export async function generateContentStream(
+  apiSettings: ApiSettings,
+  model: string,
+  params: GenerateContentRequest,
+  chromeAdapter: ChromeAdapter,
+  requestOptions?: RequestOptions
+): Promise<GenerateContentStreamResult> {
+  let response;
+  if (await chromeAdapter.isAvailable(params)) {
+    response = await chromeAdapter.generateContentStream(params);
+  } else {
+    response = await generateContentStreamOnCloud(
+      apiSettings,
+      model,
+      params,
+      requestOptions
+    );
+  }
   return processStream(response, apiSettings); // TODO: Map streaming responses
 }
 
-export async function generateContent(
+async function generateContentOnCloud(
   apiSettings: ApiSettings,
   model: string,
   params: GenerateContentRequest,
   requestOptions?: RequestOptions
-): Promise<GenerateContentResult> {
+): Promise<Response> {
   if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
     params = GoogleAIMapper.mapGenerateContentRequest(params);
   }
-  const response = await makeRequest(
+  return makeRequest(
     model,
     Task.GENERATE_CONTENT,
     apiSettings,
     /* stream */ false,
     JSON.stringify(params),
     requestOptions
   );
+}
+
+export async function generateContent(
+  apiSettings: ApiSettings,
+  model: string,
+  params: GenerateContentRequest,
+  chromeAdapter: ChromeAdapter,
+  requestOptions?: RequestOptions
+): Promise<GenerateContentResult> {
+  let response;
+  if (await chromeAdapter.isAvailable(params)) {
+    response = await chromeAdapter.generateContent(params);
+  } else {
+    response = await generateContentOnCloud(
+      apiSettings,
+      model,
+      params,
+      requestOptions
+    );
+  }
   const generateContentResponse = await processGenerateContentResponse(
     response,
     apiSettings

packages/vertexai/src/methods/generate-content.ts

@@ -22,6 +22,7 @@ import { match, restore, stub } from 'sinon';
 import { getMockResponse } from '../../test-utils/mock-response';
 import sinonChai from 'sinon-chai';
 import { VertexAIBackend } from '../backend';
+import { ChromeAdapter } from '../methods/chrome-adapter';
 
 use(sinonChai);
 
@@ -41,21 +42,27 @@ const fakeAI: AI = {
 
 describe('GenerativeModel', () => {
   it('passes params through to generateContent', async () => {
-    const genModel = new GenerativeModel(fakeAI, {
-      model: 'my-model',
-      tools: [
-        {
-          functionDeclarations: [
-            {
-              name: 'myfunc',
-              description: 'mydesc'
-            }
-          ]
-        }
-      ],
-      toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
-      systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
-    });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      {
+        model: 'my-model',
+        tools: [
+          {
+            functionDeclarations: [
+              {
+                name: 'myfunc',
+                description: 'mydesc'
+              }
+            ]
+          }
+        ],
+        toolConfig: {
+          functionCallingConfig: { mode: FunctionCallingMode.NONE }
+        },
+        systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
+      },
+      new ChromeAdapter()
+    );
     expect(genModel.tools?.length).to.equal(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
       FunctionCallingMode.NONE
@@ -86,10 +93,14 @@ describe('GenerativeModel', () => {
     restore();
   });
   it('passes text-only systemInstruction through to generateContent', async () => {
-    const genModel = new GenerativeModel(fakeAI, {
-      model: 'my-model',
-      systemInstruction: 'be friendly'
-    });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      {
+        model: 'my-model',
+        systemInstruction: 'be friendly'
+      },
+      new ChromeAdapter()
+    );
     expect(genModel.systemInstruction?.parts[0].text).to.equal('be friendly');
     const mockResponse = getMockResponse(
       'vertexAI',
@@ -112,21 +123,27 @@ describe('GenerativeModel', () => {
     restore();
   });
   it('generateContent overrides model values', async () => {
-    const genModel = new GenerativeModel(fakeAI, {
-      model: 'my-model',
-      tools: [
-        {
-          functionDeclarations: [
-            {
-              name: 'myfunc',
-              description: 'mydesc'
-            }
-          ]
-        }
-      ],
-      toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
-      systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
-    });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      {
+        model: 'my-model',
+        tools: [
+          {
+            functionDeclarations: [
+              {
+                name: 'myfunc',
+                description: 'mydesc'
+              }
+            ]
+          }
+        ],
+        toolConfig: {
+          functionCallingConfig: { mode: FunctionCallingMode.NONE }
+        },
+        systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
+      },
+      new ChromeAdapter()
+    );
     expect(genModel.tools?.length).to.equal(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
       FunctionCallingMode.NONE
@@ -168,14 +185,20 @@ describe('GenerativeModel', () => {
     restore();
   });
   it('passes params through to chat.sendMessage', async () => {
-    const genModel = new GenerativeModel(fakeAI, {
-      model: 'my-model',
-      tools: [
-        { functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
-      ],
-      toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
-      systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
-    });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      {
+        model: 'my-model',
+        tools: [
+          { functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
+        ],
+        toolConfig: {
+          functionCallingConfig: { mode: FunctionCallingMode.NONE }
+        },
+        systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
+      },
+      new ChromeAdapter()
+    );
     expect(genModel.tools?.length).to.equal(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
       FunctionCallingMode.NONE
@@ -206,10 +229,14 @@ describe('GenerativeModel', () => {
     restore();
   });
   it('passes text-only systemInstruction through to chat.sendMessage', async () => {
-    const genModel = new GenerativeModel(fakeAI, {
-      model: 'my-model',
-      systemInstruction: 'be friendly'
-    });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      {
+        model: 'my-model',
+        systemInstruction: 'be friendly'
+      },
+      new ChromeAdapter()
+    );
     expect(genModel.systemInstruction?.parts[0].text).to.equal('be friendly');
     const mockResponse = getMockResponse(
       'vertexAI',
@@ -232,14 +259,20 @@ describe('GenerativeModel', () => {
     restore();
   });
   it('startChat overrides model values', async () => {
-    const genModel = new GenerativeModel(fakeAI, {
-      model: 'my-model',
-      tools: [
-        { functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
-      ],
-      toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.NONE } },
-      systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
-    });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      {
+        model: 'my-model',
+        tools: [
+          { functionDeclarations: [{ name: 'myfunc', description: 'mydesc' }] }
+        ],
+        toolConfig: {
+          functionCallingConfig: { mode: FunctionCallingMode.NONE }
+        },
+        systemInstruction: { role: 'system', parts: [{ text: 'be friendly' }] }
+      },
+      new ChromeAdapter()
+    );
     expect(genModel.tools?.length).to.equal(1);
     expect(genModel.toolConfig?.functionCallingConfig?.mode).to.equal(
       FunctionCallingMode.NONE
@@ -284,7 +317,11 @@ describe('GenerativeModel', () => {
     restore();
   });
   it('calls countTokens', async () => {
-    const genModel = new GenerativeModel(fakeAI, { model: 'my-model' });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      { model: 'my-model' },
+      new ChromeAdapter()
+    );
     const mockResponse = getMockResponse(
       'vertexAI',
       'unary-success-total-tokens.json'

packages/vertexai/src/models/generative-model.test.ts

@@ -43,12 +43,17 @@ import {
 } from '../requests/request-helpers';
 import { AI } from '../public-types';
 import { AIModel } from './genai-model';
+import { ChromeAdapter } from '../methods/chrome-adapter';
 
 /**
  * Class for generative model APIs.
  * @public
  */
 export class GenerativeModel extends AIModel {
+  /**
+   * Defines the name of the default in-cloud model to use for hybrid inference.
+   */
+  static DEFAULT_HYBRID_IN_CLOUD_MODEL = 'gemini-2.0-flash-lite';
   generationConfig: GenerationConfig;
   safetySettings: SafetySetting[];
   requestOptions?: RequestOptions;
@@ -59,6 +64,7 @@ export class GenerativeModel extends AIModel {
   constructor(
     ai: AI,
     modelParams: ModelParams,
+    private chromeAdapter: ChromeAdapter,
     requestOptions?: RequestOptions
   ) {
     super(ai, modelParams.model);
@@ -91,6 +97,7 @@ export class GenerativeModel extends AIModel {
         systemInstruction: this.systemInstruction,
         ...formattedParams
       },
+      this.chromeAdapter,
       this.requestOptions
     );
   }
@@ -116,6 +123,7 @@ export class GenerativeModel extends AIModel {
         systemInstruction: this.systemInstruction,
         ...formattedParams
       },
+      this.chromeAdapter,
       this.requestOptions
     );
   }
@@ -128,6 +136,7 @@ export class GenerativeModel extends AIModel {
     return new ChatSession(
       this._apiSettings,
       this.model,
+      this.chromeAdapter,
       {
         tools: this.tools,
         toolConfig: this.toolConfig,
@@ -145,6 +154,11 @@ export class GenerativeModel extends AIModel {
     request: CountTokensRequest | string | Array<string | Part>
   ): Promise<CountTokensResponse> {
     const formattedParams = formatGenerateContentInput(request);
-    return countTokens(this._apiSettings, this.model, formattedParams);
+    return countTokens(
+      this._apiSettings,
+      this.model,
+      formattedParams,
+      this.chromeAdapter
+    );
   }
 }

packages/vertexai/src/models/generative-model.ts

@@ -0,0 +1,82 @@
+/**
+ * @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.
+ */
+
+export interface LanguageModel extends EventTarget {
+  create(options?: LanguageModelCreateOptions): Promise<LanguageModel>;
+  availability(options?: LanguageModelCreateCoreOptions): Promise<Availability>;
+  prompt(
+    input: LanguageModelPrompt,
+    options?: LanguageModelPromptOptions
+  ): Promise<string>;
+  promptStreaming(
+    input: LanguageModelPrompt,
+    options?: LanguageModelPromptOptions
+  ): ReadableStream;
+  measureInputUsage(
+    input: LanguageModelPrompt,
+    options?: LanguageModelPromptOptions
+  ): Promise<number>;
+  destroy(): undefined;
+}
+export enum Availability {
+  'unavailable' = 'unavailable',
+  'downloadable' = 'downloadable',
+  'downloading' = 'downloading',
+  'available' = 'available'
+}
+export interface LanguageModelCreateCoreOptions {
+  topK?: number;
+  temperature?: number;
+  expectedInputs?: LanguageModelExpectedInput[];
+}
+export interface LanguageModelCreateOptions
+  extends LanguageModelCreateCoreOptions {
+  signal?: AbortSignal;
+  systemPrompt?: string;
+  initialPrompts?: LanguageModelInitialPrompts;
+}
+interface LanguageModelPromptOptions {
+  signal?: AbortSignal;
+}
+interface LanguageModelExpectedInput {
+  type: LanguageModelMessageType;
+  languages?: string[];
+}
+// TODO: revert to type from Prompt API explainer once it's supported.
+export type LanguageModelPrompt = LanguageModelMessageContent[];
+type LanguageModelInitialPrompts =
+  | LanguageModelMessage[]
+  | LanguageModelMessageShorthand[];
+interface LanguageModelMessage {
+  role: LanguageModelMessageRole;
+  content: LanguageModelMessageContent[];
+}
+interface LanguageModelMessageShorthand {
+  role: LanguageModelMessageRole;
+  content: string;
+}
+export interface LanguageModelMessageContent {
+  type: LanguageModelMessageType;
+  content: LanguageModelMessageContentValue;
+}
+type LanguageModelMessageRole = 'system' | 'user' | 'assistant';
+type LanguageModelMessageType = 'text' | 'image' | 'audio';
+type LanguageModelMessageContentValue =
+  | ImageBitmapSource
+  | AudioBuffer
+  | BufferSource
+  | string;

packages/vertexai/src/types/language-model.ts

@@ -17,6 +17,7 @@
 
 import { TypedSchema } from '../requests/schema-builder';
 import { Content, Part } from './content';
+import { LanguageModelCreateOptions } from './language-model';
 import {
   FunctionCallingMode,
   HarmBlockMethod,
@@ -218,3 +219,29 @@ export interface FunctionCallingConfig {
   mode?: FunctionCallingMode;
   allowedFunctionNames?: string[];
 }
+
+/**
+ * Toggles hybrid inference.
+ */
+export interface HybridParams {
+  /**
+   * Specifies on-device or in-cloud inference. Defaults to prefer on-device.
+   */
+  mode: InferenceMode;
+  /**
+   * Optional. Specifies advanced params for on-device inference.
+   */
+  onDeviceParams?: LanguageModelCreateOptions;
+  /**
+   * Optional. Specifies advanced params for in-cloud inference.
+   */
+  inCloudParams?: ModelParams;
+}
+
+/**
+ * Determines whether inference happens on-device or in-cloud.
+ */
+export type InferenceMode =
+  | 'prefer_on_device'
+  | 'only_on_device'
+  | 'only_in_cloud';