Hybrid inference code changes

Author: dlarocque

packages/ai/src/api.test.ts

@@ -102,6 +102,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/ai/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';
@@ -147,16 +150,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/ai/src/backwards-compatbility.test.ts

@@ -28,6 +28,7 @@ import {
 } from './api';
 import { AI, VertexAI, AIErrorCode } from './public-types';
 import { VertexAIBackend } from './backend';
+import { ChromeAdapter } from './methods/chrome-adapter';
 
 function assertAssignable<T, _U extends T>(): void {}
 
@@ -65,7 +66,11 @@ describe('backwards-compatible types', () => {
   it('AIModel is backwards compatible with VertexAIModel', () => {
     assertAssignable<typeof VertexAIModel, typeof AIModel>();
 
-    const model = new GenerativeModel(fakeAI, { model: 'model-name' });
+    const model = new GenerativeModel(
+      fakeAI,
+      { model: 'model-name' },
+      new ChromeAdapter()
+    );
     expect(model).to.be.instanceOf(AIModel);
     expect(model).to.be.instanceOf(VertexAIModel);
   });

packages/ai/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/ai/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
     );