Adding hybrid count token implementation - First Pass

Author: gsiddh

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

@@ -18,6 +18,7 @@
 import { isChrome } from '@firebase/util';
 import {
   Content,
+  CountTokensRequest,
   GenerateContentRequest,
   InferenceMode,
   Role
@@ -102,6 +103,18 @@ export class ChromeAdapter {
     const stream = await session.promptStreaming(prompt.content);
     return ChromeAdapter.toStreamResponse(stream);
   }
+  async countTokens(request: CountTokensRequest): Promise<Response> {
+    const options = this.onDeviceParams || {};
+    const prompts = ChromeAdapter.toInitialPrompts(request.contents);
+    const session = await this.session(options);
+    const tokenCount = await session.countPromptTokens(prompts);
+    return {
+      json: async () => ({
+        totalTokens: tokenCount,
+        totalBillableCharacters: 0,
+      })
+    } as Response;
+  }
   private static isOnDeviceRequest(request: GenerateContentRequest): boolean {
     // Returns false if the prompt is empty.
     if (request.contents.length === 0) {

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

@@ -25,6 +25,7 @@ import { countTokens } from './count-tokens';
 import { CountTokensRequest } from '../types';
 import { ApiSettings } from '../types/internal';
 import { Task } from '../requests/request';
+import { ChromeAdapter } from './chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -52,7 +53,8 @@ describe('countTokens()', () => {
     const result = await countTokens(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.totalTokens).to.equal(6);
     expect(result.totalBillableCharacters).to.equal(16);
@@ -77,7 +79,8 @@ describe('countTokens()', () => {
     const result = await countTokens(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.totalTokens).to.equal(1837);
     expect(result.totalBillableCharacters).to.equal(117);
@@ -104,7 +107,8 @@ describe('countTokens()', () => {
     const result = await countTokens(
       fakeApiSettings,
       'model',
-      fakeRequestParams
+      fakeRequestParams,
+      new ChromeAdapter()
     );
     expect(result.totalTokens).to.equal(258);
     expect(result).to.not.have.property('totalBillableCharacters');
@@ -127,7 +131,7 @@ describe('countTokens()', () => {
       json: mockResponse.json
     } as Response);
     await expect(
-      countTokens(fakeApiSettings, 'model', fakeRequestParams)
+      countTokens(fakeApiSettings, 'model', fakeRequestParams, new ChromeAdapter())
     ).to.be.rejectedWith(/404.*not found/);
     expect(mockFetch).to.be.called;
   });

packages/vertexai/src/methods/count-tokens.ts

@@ -22,8 +22,9 @@ import {
 } from '../types';
 import { Task, makeRequest } from '../requests/request';
 import { ApiSettings } from '../types/internal';
+import { ChromeAdapter } from './chrome-adapter';
 
-export async function countTokens(
+export async function countTokensOnCloud(
   apiSettings: ApiSettings,
   model: string,
   params: CountTokensRequest,
@@ -39,3 +40,17 @@ export async function countTokens(
   );
   return response.json();
 }
+
+export async function countTokens(
+  apiSettings: ApiSettings,
+  model: string,
+  params: CountTokensRequest,
+  chromeAdapter: ChromeAdapter,
+  requestOptions?: RequestOptions
+): Promise<CountTokensResponse> {
+  if (await chromeAdapter.isAvailable(params)) {
+    return (await chromeAdapter.countTokens(params)).json();
+  } else {
+    return countTokensOnCloud(apiSettings, model, params, requestOptions);
+  }
+}

packages/vertexai/src/models/generative-model.ts

@@ -154,6 +154,6 @@ export class GenerativeModel extends VertexAIModel {
     request: CountTokensRequest | string | Array<string | Part>
   ): Promise<CountTokensResponse> {
     const formattedParams = formatGenerateContentInput(request);
-    return countTokens(this._apiSettings, this.model, formattedParams);
+    return countTokens(this._apiSettings, this.model, formattedParams, this.chromeAdapter);
   }
 }