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

docs-devsite/vertexai.backend.md

@@ -0,0 +1,57 @@
+Project: /docs/reference/js/_project.yaml
+Book: /docs/reference/_book.yaml
+page_type: reference
+
+{% comment %}
+DO NOT EDIT THIS FILE!
+This is generated by the JS SDK team, and any local changes will be
+overwritten. Changes should be made in the source code at
+https://github.com/firebase/firebase-js-sdk
+{% endcomment %}
+
+# Backend class
+Abstract base class representing the configuration for an AI service backend. This class should not be instantiated directly. Use its subclasses [GoogleAIBackend](./vertexai.googleaibackend.md#googleaibackend_class) or [VertexAIBackend](./vertexai.vertexaibackend.md#vertexaibackend_class)<!-- -->.
+
+<b>Signature:</b>
+
+```typescript
+export declare abstract class Backend 
+```
+
+## Constructors
+
+|  Constructor | Modifiers | Description |
+|  --- | --- | --- |
+|  [(constructor)(type)](./vertexai.backend.md#backendconstructor) |  | Protected constructor for use by subclasses. |
+
+## Properties
+
+|  Property | Modifiers | Type | Description |
+|  --- | --- | --- | --- |
+|  [backendType](./vertexai.backend.md#backendbackendtype) |  | [BackendType](./vertexai.md#backendtype) | Specifies the backend type (either 'GOOGLE\_AI' or 'VERTEX\_AI'). |
+
+## Backend.(constructor)
+
+Protected constructor for use by subclasses.
+
+<b>Signature:</b>
+
+```typescript
+protected constructor(type: BackendType);
+```
+
+#### Parameters
+
+|  Parameter | Type | Description |
+|  --- | --- | --- |
+|  type | [BackendType](./vertexai.md#backendtype) | The specific backend type constant (e.g., BackendType.GOOGLE\_AI). |
+
+## Backend.backendType
+
+Specifies the backend type (either 'GOOGLE\_AI' or 'VERTEX\_AI').
+
+<b>Signature:</b>
+
+```typescript
+readonly backendType: BackendType;
+```

docs-devsite/vertexai.googleaibackend.md

@@ -0,0 +1,36 @@
+Project: /docs/reference/js/_project.yaml
+Book: /docs/reference/_book.yaml
+page_type: reference
+
+{% comment %}
+DO NOT EDIT THIS FILE!
+This is generated by the JS SDK team, and any local changes will be
+overwritten. Changes should be made in the source code at
+https://github.com/firebase/firebase-js-sdk
+{% endcomment %}
+
+# GoogleAIBackend class
+Represents the configuration class for the Google AI backend. Use this with [AIOptions](./vertexai.aioptions.md#aioptions_interface) when initializing the service with [getAI()](./vertexai.md#getai_a94a413)<!-- -->.
+
+<b>Signature:</b>
+
+```typescript
+export declare class GoogleAIBackend extends Backend 
+```
+<b>Extends:</b> [Backend](./vertexai.backend.md#backend_class)
+
+## Constructors
+
+|  Constructor | Modifiers | Description |
+|  --- | --- | --- |
+|  [(constructor)()](./vertexai.googleaibackend.md#googleaibackendconstructor) |  | Creates a configuration object for the Google AI backend. |
+
+## GoogleAIBackend.(constructor)
+
+Creates a configuration object for the Google AI backend.
+
+<b>Signature:</b>
+
+```typescript
+constructor();
+```

docs-devsite/vertexai.hybridparams.md

@@ -0,0 +1,57 @@
+Project: /docs/reference/js/_project.yaml
+Book: /docs/reference/_book.yaml
+page_type: reference
+
+{% comment %}
+DO NOT EDIT THIS FILE!
+This is generated by the JS SDK team, and any local changes will be
+overwritten. Changes should be made in the source code at
+https://github.com/firebase/firebase-js-sdk
+{% endcomment %}
+
+# HybridParams interface
+Toggles hybrid inference.
+
+<b>Signature:</b>
+
+```typescript
+export interface HybridParams 
+```
+
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [inCloudParams](./vertexai.hybridparams.md#hybridparamsincloudparams) | [ModelParams](./vertexai.modelparams.md#modelparams_interface) | Optional. Specifies advanced params for in-cloud inference. |
+|  [mode](./vertexai.hybridparams.md#hybridparamsmode) | [InferenceMode](./vertexai.md#inferencemode) | Specifies on-device or in-cloud inference. Defaults to prefer on-device. |
+|  [onDeviceParams](./vertexai.hybridparams.md#hybridparamsondeviceparams) | LanguageModelCreateOptions | Optional. Specifies advanced params for on-device inference. |
+
+## HybridParams.inCloudParams
+
+Optional. Specifies advanced params for in-cloud inference.
+
+<b>Signature:</b>
+
+```typescript
+inCloudParams?: ModelParams;
+```
+
+## HybridParams.mode
+
+Specifies on-device or in-cloud inference. Defaults to prefer on-device.
+
+<b>Signature:</b>
+
+```typescript
+mode: InferenceMode;
+```
+
+## HybridParams.onDeviceParams
+
+Optional. Specifies advanced params for on-device inference.
+
+<b>Signature:</b>
+
+```typescript
+onDeviceParams?: LanguageModelCreateOptions;
+```

docs-devsite/vertexai.vertexaibackend.md

@@ -0,0 +1,58 @@
+Project: /docs/reference/js/_project.yaml
+Book: /docs/reference/_book.yaml
+page_type: reference
+
+{% comment %}
+DO NOT EDIT THIS FILE!
+This is generated by the JS SDK team, and any local changes will be
+overwritten. Changes should be made in the source code at
+https://github.com/firebase/firebase-js-sdk
+{% endcomment %}
+
+# VertexAIBackend class
+Represents the configuration class for the Vertex AI backend. Use this with [AIOptions](./vertexai.aioptions.md#aioptions_interface) when initializing the server with [getAI()](./vertexai.md#getai_a94a413)<!-- -->.
+
+<b>Signature:</b>
+
+```typescript
+export declare class VertexAIBackend extends Backend 
+```
+<b>Extends:</b> [Backend](./vertexai.backend.md#backend_class)
+
+## Constructors
+
+|  Constructor | Modifiers | Description |
+|  --- | --- | --- |
+|  [(constructor)(location)](./vertexai.vertexaibackend.md#vertexaibackendconstructor) |  | Creates a configuration object for the Vertex AI backend. |
+
+## Properties
+
+|  Property | Modifiers | Type | Description |
+|  --- | --- | --- | --- |
+|  [location](./vertexai.vertexaibackend.md#vertexaibackendlocation) |  | string | The region identifier. See [Vertex AI locations](https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations) for a list of supported locations. |
+
+## VertexAIBackend.(constructor)
+
+Creates a configuration object for the Vertex AI backend.
+
+<b>Signature:</b>
+
+```typescript
+constructor(location?: string);
+```
+
+#### Parameters
+
+|  Parameter | Type | Description |
+|  --- | --- | --- |
+|  location | string | The region identifier, defaulting to <code>us-central1</code>; see [Vertex AI locations](https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations) for a list of supported locations. |
+
+## VertexAIBackend.location
+
+The region identifier. See [Vertex AI locations](https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations) for a list of supported locations.
+
+<b>Signature:</b>
+
+```typescript
+readonly location: string;
+```

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,