VertexAI: add test cases for countTokens() (#8317)

Author: tanzimfh

Diff for: packages/vertexai/src/methods/count-tokens.test.ts

@@ -0,0 +1,110 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { expect, use } from 'chai';
+import { match, restore, stub } from 'sinon';
+import sinonChai from 'sinon-chai';
+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 { ApiSettings } from '../types/internal';
+import { Task } from '../requests/request';
+
+use(sinonChai);
+use(chaiAsPromised);
+
+const fakeApiSettings: ApiSettings = {
+  apiKey: 'key',
+  project: 'my-project',
+  location: 'us-central1'
+};
+
+const fakeRequestParams: CountTokensRequest = {
+  contents: [{ parts: [{ text: 'hello' }], role: 'user' }]
+};
+
+describe('countTokens()', () => {
+  afterEach(() => {
+    restore();
+  });
+  it('total tokens', async () => {
+    const mockResponse = getMockResponse(
+      'count-tokens-success-total-tokens.json'
+    );
+    const makeRequestStub = stub(request, 'makeRequest').resolves(
+      mockResponse as Response
+    );
+    const result = await countTokens(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams
+    );
+    expect(result.totalTokens).to.equal(6);
+    expect(result.totalBillableCharacters).to.equal(16);
+    expect(makeRequestStub).to.be.calledWith(
+      'model',
+      Task.COUNT_TOKENS,
+      fakeApiSettings,
+      false,
+      match((value: string) => {
+        return value.includes('contents');
+      }),
+      undefined
+    );
+  });
+  it('total tokens no billable characters', async () => {
+    const mockResponse = getMockResponse(
+      'count-tokens-success-no-billable-characters.json'
+    );
+    const makeRequestStub = stub(request, 'makeRequest').resolves(
+      mockResponse as Response
+    );
+    const result = await countTokens(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams
+    );
+    expect(result.totalTokens).to.equal(258);
+    expect(result).to.not.have.property('totalBillableCharacters');
+    expect(makeRequestStub).to.be.calledWith(
+      'model',
+      Task.COUNT_TOKENS,
+      fakeApiSettings,
+      false,
+      match((value: string) => {
+        return value.includes('contents');
+      }),
+      undefined
+    );
+  });
+  it('model not found', async () => {
+    const mockResponse = getMockResponse(
+      'count-tokens-failure-model-not-found.json'
+    );
+    const mockFetch = stub(globalThis, 'fetch').resolves({
+      ok: false,
+      status: 404,
+      json: mockResponse.json
+    } as Response);
+    await expect(
+      countTokens(fakeApiSettings, 'model', fakeRequestParams)
+    ).to.be.rejectedWith(/404.*not found/);
+    expect(mockFetch).to.be.called;
+  });
+});

Diff for: packages/vertexai/src/models/generative-model.test.ts

@@ -262,4 +262,24 @@ describe('GenerativeModel', () => {
     );
     restore();
   });
+  it('calls countTokens', async () => {
+    const genModel = new GenerativeModel(fakeVertexAI, { model: 'my-model' });
+    const mockResponse = getMockResponse(
+      'count-tokens-success-total-tokens.json'
+    );
+    const makeRequestStub = stub(request, 'makeRequest').resolves(
+      mockResponse as Response
+    );
+    await genModel.countTokens('hello');
+    expect(makeRequestStub).to.be.calledWith(
+      'publishers/google/models/my-model',
+      request.Task.COUNT_TOKENS,
+      match.any,
+      false,
+      match((value: string) => {
+        return value.includes('hello');
+      })
+    );
+    restore();
+  });
 });

Diff for: packages/vertexai/test-utils/mock-responses/count-tokens-failure-model-not-found.json

@@ -0,0 +1,13 @@
+{
+  "error": {
+    "code": 404,
+    "message": "models/test-model-name is not found for API version v1beta, or is not supported for countTokens. Call ListModels to see the list of available models and their supported methods.",
+    "status": "NOT_FOUND",
+    "details": [
+      {
+        "@type": "type.googleapis.com/google.rpc.DebugInfo",
+        "detail": "[ORIGINAL ERROR] generic::not_found: models/test-model-name is not found for API version v1beta, or is not supported for countTokens. Call ListModels to see the list of available models and their supported methods. [google.rpc.error_details_ext] { message: \"models/test-model-name is not found for API version v1beta, or is not supported for countTokens. Call ListModels to see the list of available models and their supported methods.\" }"
+      }
+    ]
+  }
+}

Diff for: packages/vertexai/test-utils/mock-responses/count-tokens-success-no-billable-characters.json

@@ -0,0 +1,3 @@
+{
+  "totalTokens": 258
+}

Diff for: packages/vertexai/test-utils/mock-responses/count-tokens-success-total-tokens.json

@@ -0,0 +1,4 @@
+{
+  "totalTokens": 6,
+  "totalBillableCharacters": 16
+}