feat: add countTokens method

PiperOrigin-RevId: 586023998

Author: sararob

src/index.ts

@@ -19,7 +19,7 @@
 import {GoogleAuth} from 'google-auth-library';
 
 import {emptyGenerator, processStream} from './process_stream';
-import {Content, GenerateContentParams, GenerateContentRequest, GenerateContentResult, GenerationConfig, ModelParams, Part, SafetySetting} from './types/content';
+import {Content, CountTokensRequest, CountTokensResponse, GenerateContentParams, GenerateContentRequest, GenerateContentResult, GenerationConfig, ModelParams, Part, SafetySetting} from './types/content';
 import {postRequest} from './util';
 
 // TODO: update this when model names are available
@@ -264,6 +264,42 @@ export class GenerativeModel {
     }
   }
 
+  /**
+   * Make a countTokens request.
+   * @param request A CountTokensRequest object with the request contents.
+   * @return The CountTokensResponse object with the token count.
+   */
+  async countTokens(request: CountTokensRequest): Promise<CountTokensResponse> {
+    let response;
+    try {
+      response = await postRequest({
+        region: this._vertex_instance.location,
+        project: this._vertex_instance.project,
+        resourcePath: `publishers/google/models/${this.model}`,
+        resourceMethod: 'countTokens',
+        token: await this._vertex_instance.token,
+        data: request,
+        apiEndpoint: this._vertex_instance.apiEndpoint,
+      });
+      if (response === undefined) {
+        throw new Error('did not get a valid response.');
+      }
+      if (!response.ok) {
+        throw new Error(`${response.status} ${response.statusText}`);
+      }
+
+    } catch (e) {
+      console.log(e);
+    }
+    if (response) {
+      const responseJson = await response.json();
+      return responseJson as CountTokensResponse;
+    } else {
+      throw new Error('did not get a valid response.');
+    }
+  }
+
+
   startChat(request: StartChatParams): ChatSession {
     const startChatRequest = {
       history: request.history,

src/index_test.ts

@@ -20,7 +20,7 @@ import 'jasmine';
 
 import {ChatSession, GenerativeModel, StartChatParams, VertexAI} from './index';
 import * as StreamFunctions from './process_stream';
-import {GenerateContentParams, GenerateContentResult} from './types/content';
+import {CountTokensRequest, CountTokensResponse, GenerateContentParams, GenerateContentResult} from './types/content';
 import * as PostRequest from './util/post_request';
 
 const PROJECT = 'test_project';
@@ -135,6 +135,26 @@ describe('VertexAI', () => {
       expect(resp).toBeInstanceOf(ChatSession);
     });
   });
+
+  describe('countTokens', () => {
+    it('returns the token count', async () => {
+      const req: CountTokensRequest = {
+        contents: TEST_USER_CHAT_MESSAGE,
+      };
+      const responseBody = {
+        totalTokens: 1,
+      };
+      const response = new Response(JSON.stringify(responseBody), {
+        status: 200,
+        statusText: 'OK',
+        headers: {'Content-Type': 'application/json'},
+      });
+      const responsePromise = Promise.resolve(response);
+      spyOn(global, 'fetch').and.returnValue(responsePromise);
+      const resp = await model.countTokens(req);
+      expect(resp).toEqual(responseBody);
+    });
+  });
 });
 
 describe('ChatSession', () => {

src/types/content.ts

@@ -31,6 +31,21 @@ export declare interface GenerateContentRequest extends BaseModelParams {
   contents: Content[];
 }
 
+/**
+ * Params used to call countTokens
+ */
+export declare interface CountTokensRequest {
+  contents: Content[];
+}
+
+/**
+ * Response returned from countTokens
+ */
+export declare interface CountTokensResponse {
+  totalTokens: number;
+  totalBillableCharacters?: number;
+}
+
 
 /**
  * Configuration for initializing a model, for example via getGenerativeModel

src/util/post_request.ts

@@ -18,9 +18,7 @@
 // TODO: update to prod endpoint when ready
 const API_BASE_PATH = 'autopush-aiplatform.sandbox.googleapis.com';
 
-import {
-    GenerateContentRequest, CLIENT_INFO,
-} from '../types/content';
+import {GenerateContentRequest, CLIENT_INFO, CountTokensRequest} from '../types/content';
 
 /**
  * Makes a POST request to a Vertex service
@@ -41,7 +39,7 @@ export async function postRequest({
   resourcePath: string,
   resourceMethod: string,
   token: string,
-  data: GenerateContentRequest,
+  data: GenerateContentRequest|CountTokensRequest,
   apiEndpoint?: string,
   apiVersion?: string,
 }): Promise<Response|undefined> {