Define ChromeAdapter class (#8942)

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

Author: gsiddh

common/api-review/vertexai.api.md

@@ -42,7 +42,8 @@ export class BooleanSchema extends Schema {
 // @public
 export class ChatSession {
     // Warning: (ae-forgotten-export) The symbol "ApiSettings" needs to be exported by the entry point index.d.ts
-    constructor(apiSettings: ApiSettings, model: string, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
+    // Warning: (ae-forgotten-export) The symbol "ChromeAdapter" needs to be exported by the entry point index.d.ts
+    constructor(apiSettings: ApiSettings, model: string, chromeAdapter: ChromeAdapter, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
     getHistory(): Promise<Content[]>;
     // (undocumented)
     model: string;
@@ -324,7 +325,7 @@ export interface GenerativeContentBlob {
 
 // @public
 export class GenerativeModel extends VertexAIModel {
-    constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions);
+    constructor(vertexAI: VertexAI, modelParams: ModelParams, chromeAdapter: ChromeAdapter, requestOptions?: RequestOptions);
     countTokens(request: CountTokensRequest | string | Array<string | Part>): Promise<CountTokensResponse>;
     static DEFAULT_HYBRID_IN_CLOUD_MODEL: string;
     generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;

docs-devsite/vertexai.chatsession.md

@@ -22,7 +22,7 @@ export declare class ChatSession
 
 |  Constructor | Modifiers | Description |
 |  --- | --- | --- |
-|  [(constructor)(apiSettings, model, params, requestOptions)](./vertexai.chatsession.md#chatsessionconstructor) |  | Constructs a new instance of the <code>ChatSession</code> class |
+|  [(constructor)(apiSettings, model, chromeAdapter, params, requestOptions)](./vertexai.chatsession.md#chatsessionconstructor) |  | Constructs a new instance of the <code>ChatSession</code> class |
 
 ## Properties
 
@@ -47,7 +47,7 @@ Constructs a new instance of the `ChatSession` class
 <b>Signature:</b>
 
 ```typescript
-constructor(apiSettings: ApiSettings, model: string, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
+constructor(apiSettings: ApiSettings, model: string, chromeAdapter: ChromeAdapter, params?: StartChatParams | undefined, requestOptions?: RequestOptions | undefined);
 ```
 
 #### Parameters
@@ -56,6 +56,7 @@ constructor(apiSettings: ApiSettings, model: string, params?: StartChatParams |
 |  --- | --- | --- |
 |  apiSettings | ApiSettings |  |
 |  model | string |  |
+|  chromeAdapter | ChromeAdapter |  |
 |  params | [StartChatParams](./vertexai.startchatparams.md#startchatparams_interface) \| undefined |  |
 |  requestOptions | [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) \| undefined |  |
 

docs-devsite/vertexai.generativemodel.md

@@ -23,7 +23,7 @@ export declare class GenerativeModel extends VertexAIModel
 
 |  Constructor | Modifiers | Description |
 |  --- | --- | --- |
-|  [(constructor)(vertexAI, modelParams, requestOptions)](./vertexai.generativemodel.md#generativemodelconstructor) |  | Constructs a new instance of the <code>GenerativeModel</code> class |
+|  [(constructor)(vertexAI, modelParams, chromeAdapter, requestOptions)](./vertexai.generativemodel.md#generativemodelconstructor) |  | Constructs a new instance of the <code>GenerativeModel</code> class |
 
 ## Properties
 
@@ -53,7 +53,7 @@ Constructs a new instance of the `GenerativeModel` class
 <b>Signature:</b>
 
 ```typescript
-constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions);
+constructor(vertexAI: VertexAI, modelParams: ModelParams, chromeAdapter: ChromeAdapter, requestOptions?: RequestOptions);
 ```
 
 #### Parameters
@@ -62,6 +62,7 @@ constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: Reque
 |  --- | --- | --- |
 |  vertexAI | [VertexAI](./vertexai.vertexai.md#vertexai_interface) |  |
 |  modelParams | [ModelParams](./vertexai.modelparams.md#modelparams_interface) |  |
+|  chromeAdapter | ChromeAdapter |  |
 |  requestOptions | [RequestOptions](./vertexai.requestoptions.md#requestoptions_interface) |  |
 
 ## GenerativeModel.DEFAULT\_HYBRID\_IN\_CLOUD\_MODEL

packages/vertexai/src/api.ts

@@ -30,6 +30,7 @@ import {
 } from './types';
 import { VertexAIError } from './errors';
 import { VertexAIModel, GenerativeModel, ImagenModel } from './models';
+import { ChromeAdapter } from './methods/chrome-adapter';
 
 export { ChatSession } from './methods/chat-session';
 export * from './requests/schema-builder';
@@ -91,7 +92,12 @@ export function getGenerativeModel(
       `Must provide a model name. Example: getGenerativeModel({ model: 'my-model-name' })`
     );
   }
-  return new GenerativeModel(vertexAI, inCloudParams, requestOptions);
+  return new GenerativeModel(
+    vertexAI,
+    inCloudParams,
+    new ChromeAdapter(hybridParams.mode, hybridParams.onDeviceParams),
+    requestOptions
+  );
 }
 
 /**

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

@@ -23,6 +23,7 @@ import * as generateContentMethods from './generate-content';
 import { GenerateContentStreamResult } from '../types';
 import { ChatSession } from './chat-session';
 import { ApiSettings } from '../types/internal';
+import { ChromeAdapter } from './chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -44,7 +45,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,
@@ -61,7 +66,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,
@@ -80,7 +89,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
         )
       )

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

@@ -0,0 +1,51 @@
+/**
+ * @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 { GenerateContentRequest, InferenceMode } from '../types';
+import { LanguageModelCreateOptions } 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 {
+  constructor(
+    private mode?: InferenceMode,
+    private onDeviceParams?: LanguageModelCreateOptions
+  ) {}
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  async isAvailable(request: GenerateContentRequest): Promise<boolean> {
+    return false;
+  }
+  async generateContentOnDevice(
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    request: GenerateContentRequest
+  ): Promise<Response> {
+    return {
+      json: () =>
+        Promise.resolve({
+          candidates: [
+            {
+              content: {
+                parts: [{ text: '' }]
+              }
+            }
+          ]
+        })
+    } as Response;
+  }
+}

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

@@ -30,6 +30,7 @@ import {
 } from '../types';
 import { ApiSettings } from '../types/internal';
 import { Task } from '../requests/request';
+import { ChromeAdapter } from './chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -70,7 +71,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(
@@ -95,7 +97,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');
@@ -118,7 +121,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);
@@ -153,7 +157,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'
@@ -180,7 +185,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(
@@ -202,7 +208,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(
@@ -224,7 +231,8 @@ describe('generateContent()', () => {
     const result = await generateContent(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.response.text()).to.equal('');
     expect(makeRequestStub).to.be.calledWith(
@@ -246,7 +254,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(
@@ -268,7 +277,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;
   });
@@ -283,10 +297,36 @@ 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/
     );
     expect(mockFetch).to.be.called;
   });
+  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,
+      'generateContentOnDevice'
+    ).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);
+  });
 });