change back to private

Author: hsubox76

common/api-review/ai.api.md

@@ -143,6 +143,8 @@ export interface ChromeAdapter {
     generateContent(request: GenerateContentRequest): Promise<Response>;
     generateContentStream(request: GenerateContentRequest): Promise<Response>;
     isAvailable(request: GenerateContentRequest): Promise<boolean>;
+    // (undocumented)
+    mode: InferenceMode;
 }
 
 // @public

docs-devsite/ai.chromeadapter.md

@@ -22,6 +22,12 @@ These methods should not be called directly by the user.
 export interface ChromeAdapter 
 ```
 
+## Properties
+
+|  Property | Type | Description |
+|  --- | --- | --- |
+|  [mode](./ai.chromeadapter.md#chromeadaptermode) | [InferenceMode](./ai.md#inferencemode) |  |
+
 ## Methods
 
 |  Method | Description |
@@ -31,6 +37,14 @@ export interface ChromeAdapter
 |  [generateContentStream(request)](./ai.chromeadapter.md#chromeadaptergeneratecontentstream) | Generates content stream on device.<p>This is comparable to [GenerativeModel.generateContentStream()](./ai.generativemodel.md#generativemodelgeneratecontentstream) for generating content in Cloud.</p> |
 |  [isAvailable(request)](./ai.chromeadapter.md#chromeadapterisavailable) | 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> |
 
+## ChromeAdapter.mode
+
+<b>Signature:</b>
+
+```typescript
+mode: InferenceMode;
+```
+
 ## ChromeAdapter.countTokens()
 
 Stub - not yet available for on-device.
@@ -73,7 +87,7 @@ generateContent(request: GenerateContentRequest): Promise<Response>;
 
 Promise&lt;Response&gt;
 
-, so we can reuse common response formatting.
+Response, so we can reuse common response formatting.
 
 ## ChromeAdapter.generateContentStream()
 
@@ -97,7 +111,7 @@ generateContentStream(request: GenerateContentRequest): Promise<Response>;
 
 Promise&lt;Response&gt;
 
-, so we can reuse common response formatting.
+Response, so we can reuse common response formatting.
 
 ## ChromeAdapter.isAvailable()
 

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

@@ -20,7 +20,7 @@ import { match, restore, stub, useFakeTimers } from 'sinon';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import * as generateContentMethods from './generate-content';
-import { GenerateContentStreamResult } from '../types';
+import { GenerateContentStreamResult, InferenceMode } from '../types';
 import { ChatSession } from './chat-session';
 import { ApiSettings } from '../types/internal';
 import { VertexAIBackend } from '../backend';
@@ -37,6 +37,12 @@ const fakeApiSettings: ApiSettings = {
   backend: new VertexAIBackend()
 };
 
+const fakeChromeAdapter = new ChromeAdapterImpl(
+  // @ts-expect-error
+  undefined,
+  InferenceMode.PREFER_ON_DEVICE
+);
+
 describe('ChatSession', () => {
   afterEach(() => {
     restore();
@@ -50,7 +56,7 @@ describe('ChatSession', () => {
       const chatSession = new ChatSession(
         fakeApiSettings,
         'a-model',
-        new ChromeAdapterImpl()
+        fakeChromeAdapter
       );
       await expect(chatSession.sendMessage('hello')).to.be.rejected;
       expect(generateContentStub).to.be.calledWith(
@@ -71,7 +77,7 @@ describe('ChatSession', () => {
       const chatSession = new ChatSession(
         fakeApiSettings,
         'a-model',
-        new ChromeAdapterImpl()
+        fakeChromeAdapter
       );
       await expect(chatSession.sendMessageStream('hello')).to.be.rejected;
       expect(generateContentStreamStub).to.be.calledWith(
@@ -94,7 +100,7 @@ describe('ChatSession', () => {
       const chatSession = new ChatSession(
         fakeApiSettings,
         'a-model',
-        new ChromeAdapterImpl()
+        fakeChromeAdapter
       );
       await chatSession.sendMessageStream('hello');
       expect(generateContentStreamStub).to.be.calledWith(

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

@@ -56,12 +56,12 @@ describe('ChromeAdapter', () => {
   describe('constructor', () => {
     it('sets image as expected input type by default', async () => {
       const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.available)
+        availability: () => Promise.resolve(Availability.AVAILABLE)
       } as LanguageModel;
       const availabilityStub = stub(
         languageModelProvider,
         'availability'
-      ).resolves(Availability.available);
+      ).resolves(Availability.AVAILABLE);
       const adapter = new ChromeAdapterImpl(
         languageModelProvider,
         InferenceMode.PREFER_ON_DEVICE
@@ -80,12 +80,12 @@ describe('ChromeAdapter', () => {
     });
     it('honors explicitly set expected inputs', async () => {
       const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.available)
+        availability: () => Promise.resolve(Availability.AVAILABLE)
       } as LanguageModel;
       const availabilityStub = stub(
         languageModelProvider,
         'availability'
-      ).resolves(Availability.available);
+      ).resolves(Availability.AVAILABLE);
       const createOptions = {
         // Explicitly sets expected inputs.
         expectedInputs: [{ type: 'text' }]
@@ -109,16 +109,11 @@ describe('ChromeAdapter', () => {
     });
   });
   describe('isAvailable', () => {
-    it('returns false if mode is undefined', async () => {
-      const adapter = new ChromeAdapterImpl();
-      expect(
-        await adapter.isAvailable({
-          contents: []
-        })
-      ).to.be.false;
-    });
     it('returns false if mode is only cloud', async () => {
-      const adapter = new ChromeAdapterImpl(undefined, 'only_in_cloud');
+      const adapter = new ChromeAdapterImpl(
+        {} as LanguageModel,
+        'only_in_cloud'
+      );
       expect(
         await adapter.isAvailable({
           contents: []
@@ -127,6 +122,7 @@ describe('ChromeAdapter', () => {
     });
     it('returns false if LanguageModel API is undefined', async () => {
       const adapter = new ChromeAdapterImpl(
+        // @ts-expect-error
         undefined,
         InferenceMode.PREFER_ON_DEVICE
       );
@@ -139,7 +135,7 @@ describe('ChromeAdapter', () => {
     it('returns false if request contents empty', async () => {
       const adapter = new ChromeAdapterImpl(
         {
-          availability: async () => Availability.available
+          availability: async () => Availability.AVAILABLE
         } as LanguageModel,
         InferenceMode.PREFER_ON_DEVICE
       );
@@ -152,7 +148,7 @@ describe('ChromeAdapter', () => {
     it('returns false if request content has "function" role', async () => {
       const adapter = new ChromeAdapterImpl(
         {
-          availability: async () => Availability.available
+          availability: async () => Availability.AVAILABLE
         } as LanguageModel,
         InferenceMode.PREFER_ON_DEVICE
       );
@@ -170,7 +166,7 @@ describe('ChromeAdapter', () => {
     it('returns true if request has image with supported mime type', async () => {
       const adapter = new ChromeAdapterImpl(
         {
-          availability: async () => Availability.available
+          availability: async () => Availability.AVAILABLE
         } as LanguageModel,
         InferenceMode.PREFER_ON_DEVICE
       );
@@ -196,7 +192,7 @@ describe('ChromeAdapter', () => {
     });
     it('returns true if model is readily available', async () => {
       const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.available)
+        availability: () => Promise.resolve(Availability.AVAILABLE)
       } as LanguageModel;
       const adapter = new ChromeAdapterImpl(
         languageModelProvider,
@@ -218,7 +214,7 @@ describe('ChromeAdapter', () => {
     });
     it('returns false and triggers download when model is available after download', async () => {
       const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.downloadable),
+        availability: () => Promise.resolve(Availability.DOWNLOADABLE),
         create: () => Promise.resolve({})
       } as LanguageModel;
       const createStub = stub(languageModelProvider, 'create').resolves(
@@ -241,7 +237,7 @@ describe('ChromeAdapter', () => {
     });
     it('avoids redundant downloads', async () => {
       const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.downloadable),
+        availability: () => Promise.resolve(Availability.DOWNLOADABLE),
         create: () => Promise.resolve({})
       } as LanguageModel;
       const downloadPromise = new Promise<LanguageModel>(() => {
@@ -264,7 +260,7 @@ describe('ChromeAdapter', () => {
     });
     it('clears state when download completes', async () => {
       const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.downloadable),
+        availability: () => Promise.resolve(Availability.DOWNLOADABLE),
         create: () => Promise.resolve({})
       } as LanguageModel;
       let resolveDownload;
@@ -289,7 +285,7 @@ describe('ChromeAdapter', () => {
     });
     it('returns false when model is never available', async () => {
       const languageModelProvider = {
-        availability: () => Promise.resolve(Availability.unavailable),
+        availability: () => Promise.resolve(Availability.UNAVAILABLE),
         create: () => Promise.resolve({})
       } as LanguageModel;
       const adapter = new ChromeAdapterImpl(
@@ -305,6 +301,7 @@ describe('ChromeAdapter', () => {
   });
   describe('generateContent', () => {
     it('throws if Chrome API is undefined', async () => {
+      // @ts-expect-error
       const adapter = new ChromeAdapterImpl(undefined, 'only_on_device');
       await expect(
         adapter.generateContent({

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

@@ -47,8 +47,8 @@ export class ChromeAdapterImpl implements ChromeAdapter {
   private downloadPromise: Promise<LanguageModel | void> | undefined;
   private oldSession: LanguageModel | undefined;
   constructor(
-    private languageModelProvider?: LanguageModel,
-    private mode?: InferenceMode,
+    private languageModelProvider: LanguageModel,
+    private mode: InferenceMode,
     private onDeviceParams: OnDeviceParams = {
       createOptions: {
         // Defaults to support image inputs for convenience.
@@ -79,7 +79,7 @@ export class ChromeAdapterImpl implements ChromeAdapter {
       );
       return false;
     }
-    if (this.mode === 'only_in_cloud') {
+    if (this.mode === InferenceMode.ONLY_IN_CLOUD) {
       logger.debug(
         `On-device inference unavailable because mode is "only_in_cloud".`
       );
@@ -89,12 +89,27 @@ export class ChromeAdapterImpl implements ChromeAdapter {
     // Triggers out-of-band download so model will eventually become available.
     const availability = await this.downloadIfAvailable();
 
-    if (this.mode === 'only_on_device') {
+    if (this.mode === InferenceMode.ONLY_ON_DEVICE) {
+      // If it will never be available due to API inavailability, throw.
+      if (availability === Availability.UNAVAILABLE) {
+        throw new AIError(
+          AIErrorCode.API_NOT_ENABLED,
+          'Local LanguageModel API not available in this environment.'
+        );
+      } else if (
+        availability === Availability.DOWNLOADABLE ||
+        availability === Availability.DOWNLOADING
+      ) {
+        // TODO(chholland): Better user experience during download - progress?
+        logger.debug(`Waiting for download of LanguageModel to complete.`);
+        await this.downloadPromise;
+        return true;
+      }
       return true;
     }
 
     // Applies prefer_on_device logic.
-    if (availability !== Availability.available) {
+    if (availability !== Availability.AVAILABLE) {
       logger.debug(
         `On-device inference unavailable because availability is "${availability}".`
       );
@@ -202,7 +217,7 @@ export class ChromeAdapterImpl implements ChromeAdapter {
       this.onDeviceParams.createOptions
     );
 
-    if (availability === Availability.downloadable) {
+    if (availability === Availability.DOWNLOADABLE) {
       this.download();
     }
 

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

@@ -22,7 +22,7 @@ import chaiAsPromised from 'chai-as-promised';
 import { getMockResponse } from '../../test-utils/mock-response';
 import * as request from '../requests/request';
 import { countTokens } from './count-tokens';
-import { CountTokensRequest } from '../types';
+import { CountTokensRequest, InferenceMode } from '../types';
 import { ApiSettings } from '../types/internal';
 import { Task } from '../requests/request';
 import { mapCountTokensRequest } from '../googleai-mappers';
@@ -52,6 +52,12 @@ const fakeRequestParams: CountTokensRequest = {
   contents: [{ parts: [{ text: 'hello' }], role: 'user' }]
 };
 
+const fakeChromeAdapter = new ChromeAdapterImpl(
+  // @ts-expect-error
+  undefined,
+  InferenceMode.PREFER_ON_DEVICE
+);
+
 describe('countTokens()', () => {
   afterEach(() => {
     restore();
@@ -68,7 +74,7 @@ describe('countTokens()', () => {
       fakeApiSettings,
       'model',
       fakeRequestParams,
-      new ChromeAdapterImpl()
+      fakeChromeAdapter
     );
     expect(result.totalTokens).to.equal(6);
     expect(result.totalBillableCharacters).to.equal(16);
@@ -95,7 +101,7 @@ describe('countTokens()', () => {
       fakeApiSettings,
       'model',
       fakeRequestParams,
-      new ChromeAdapterImpl()
+      fakeChromeAdapter
     );
     expect(result.totalTokens).to.equal(1837);
     expect(result.totalBillableCharacters).to.equal(117);
@@ -124,7 +130,7 @@ describe('countTokens()', () => {
       fakeApiSettings,
       'model',
       fakeRequestParams,
-      new ChromeAdapterImpl()
+      fakeChromeAdapter
     );
     expect(result.totalTokens).to.equal(258);
     expect(result).to.not.have.property('totalBillableCharacters');
@@ -154,7 +160,7 @@ describe('countTokens()', () => {
         fakeApiSettings,
         'model',
         fakeRequestParams,
-        new ChromeAdapterImpl()
+        fakeChromeAdapter
       )
     ).to.be.rejectedWith(/404.*not found/);
     expect(mockFetch).to.be.called;
@@ -177,7 +183,7 @@ describe('countTokens()', () => {
         fakeGoogleAIApiSettings,
         'model',
         fakeRequestParams,
-        new ChromeAdapterImpl()
+        fakeChromeAdapter
       );
 
       expect(makeRequestStub).to.be.calledWith(
@@ -191,7 +197,7 @@ describe('countTokens()', () => {
     });
   });
   it('on-device', async () => {
-    const chromeAdapter = new ChromeAdapterImpl();
+    const chromeAdapter = fakeChromeAdapter;
     const isAvailableStub = stub(chromeAdapter, 'isAvailable').resolves(true);
     const mockResponse = getMockResponse(
       'vertexAI',