fix(ai): Add a warning when setting top-level params in hybrid mode (#9577)

Author: hsubox76

packages/ai/src/api.test.ts

@@ -14,7 +14,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { ImagenModelParams, ModelParams, AIErrorCode } from './types';
+import {
+  ImagenModelParams,
+  ModelParams,
+  AIErrorCode,
+  InferenceMode
+} from './types';
 import { AIError } from './errors';
 import {
   getAI,
@@ -28,12 +33,17 @@ import {
   getTemplateImagenModel,
   TemplateImagenModel
 } from './api';
-import { expect } from 'chai';
+import { expect, use } from 'chai';
+import { stub } from 'sinon';
 import { AI } from './public-types';
 import { GenerativeModel } from './models/generative-model';
 import { GoogleAIBackend, VertexAIBackend } from './backend';
 import { getFullApp } from '../test-utils/get-fake-firebase-services';
 import { AI_TYPE } from './constants';
+import { logger } from './logger';
+import sinonChai from 'sinon-chai';
+
+use(sinonChai);
 
 const fakeAI: AI = {
   app: {
@@ -148,6 +158,17 @@ describe('Top level API', () => {
     expect(genModel).to.be.an.instanceOf(GenerativeModel);
     expect(genModel.model).to.equal('publishers/google/models/my-model');
   });
+  it('getGenerativeModel warns in hybrid mode if top-level params are set', () => {
+    const warnStub = stub(logger, 'warn');
+    const genModel = getGenerativeModel(fakeAI, {
+      mode: InferenceMode.PREFER_ON_DEVICE,
+      generationConfig: {}
+    });
+    expect(genModel).to.be.an.instanceOf(GenerativeModel);
+    expect(warnStub).to.be.calledWithMatch(InferenceMode.PREFER_ON_DEVICE);
+    expect(warnStub).to.be.calledWithMatch('generationConfig');
+    warnStub.restore();
+  });
   it('getImagenModel throws if no model is provided', () => {
     try {
       getImagenModel(fakeAI, {} as ImagenModelParams);

packages/ai/src/api.ts

@@ -41,6 +41,7 @@ import { GoogleAIBackend } from './backend';
 import { WebSocketHandlerImpl } from './websocket';
 import { TemplateGenerativeModel } from './models/template-generative-model';
 import { TemplateImagenModel } from './models/template-imagen-model';
+import { logger } from './logger';
 
 export { ChatSession } from './methods/chat-session';
 export { LiveSession } from './methods/live-session';
@@ -117,6 +118,12 @@ export function getAI(app: FirebaseApp = getApp(), options?: AIOptions): AI {
   return aiInstance;
 }
 
+const hybridParamKeys: Array<keyof HybridParams> = [
+  'mode',
+  'onDeviceParams',
+  'inCloudParams'
+];
+
 /**
  * Returns a {@link GenerativeModel} class with methods for inference
  * and other functionality.
@@ -132,6 +139,18 @@ export function getGenerativeModel(
   const hybridParams = modelParams as HybridParams;
   let inCloudParams: ModelParams;
   if (hybridParams.mode) {
+    for (const param of Object.keys(modelParams)) {
+      if (!hybridParamKeys.includes(param as keyof HybridParams)) {
+        logger.warn(
+          `When a hybrid inference mode is specified (mode is currently set` +
+            ` to ${hybridParams.mode}), "${param}" cannot be ` +
+            `configured at the top level. Configuration for in-cloud and ` +
+            `on-device must be done separately in inCloudParams and onDeviceParams. ` +
+            `Configuration values set outside of inCloudParams and onDeviceParams will` +
+            ` be ignored.`
+        );
+      }
+    }
     inCloudParams = hybridParams.inCloudParams || {
       model: DEFAULT_HYBRID_IN_CLOUD_MODEL
     };