Add support for token-based usage metrics.

common/api-review/vertexai.api.md

@@ -92,6 +92,8 @@ export interface CountTokensRequest {
 
 // @public
 export interface CountTokensResponse {
+    // (undocumented)
+    promptTokensDetails?: ModalityTokenCount[];
     totalBillableCharacters?: number;
     totalTokens: number;
 }
@@ -447,6 +449,30 @@ export class IntegerSchema extends Schema {
     constructor(schemaParams?: SchemaParams);
 }
 
+// @public
+export enum Modality {
+    // (undocumented)
+    AUDIO = "AUDIO",
+    // (undocumented)
+    DOCUMENT = "DOCUMENT",
+    // (undocumented)
+    IMAGE = "IMAGE",
+    // (undocumented)
+    MODALITY_UNSPECIFIED = "MODALITY_UNSPECIFIED",
+    // (undocumented)
+    TEXT = "TEXT",
+    // (undocumented)
+    VIDEO = "VIDEO"
+}
+
+// @public
+export interface ModalityTokenCount {
+    // (undocumented)
+    modality: Modality;
+    // (undocumented)
+    tokenCount: number;
+}
+
 // @public
 export interface ModelParams extends BaseParams {
     // (undocumented)
@@ -682,8 +708,12 @@ export interface UsageMetadata {
     // (undocumented)
     candidatesTokenCount: number;
     // (undocumented)
+    candidatesTokensDetails?: ModalityTokenCount[];
+    // (undocumented)
     promptTokenCount: number;
     // (undocumented)
+    promptTokensDetails?: ModalityTokenCount[];
+    // (undocumented)
     totalTokenCount: number;
 }
 

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

@@ -66,6 +66,31 @@ describe('countTokens()', () => {
       undefined
     );
   });
+  it('total tokens with modality details', async () => {
+    const mockResponse = getMockResponse('unary-success-detailed-token-response.json');
+    const makeRequestStub = stub(request, 'makeRequest').resolves(
+      mockResponse as Response
+    );
+    const result = await countTokens(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams
+    );
+    expect(result.totalTokens).to.equal(1837);
+    expect(result.totalBillableCharacters).to.equal(117);
+    expect(result.promptTokensDetails?.[0].modality).to.equal('IMAGE');
+    expect(result.promptTokensDetails?.[0].tokenCount).to.equal(1806);
+    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(
       'unary-success-no-billable-characters.json'

packages/vertexai/src/methods/generate-content.test.ts

@@ -102,6 +102,30 @@ describe('generateContent()', () => {
       match.any
     );
   });
+  it('long response with token details', async () => {
+    const mockResponse = getMockResponse('unary-success-basic-response-long-usage-metadata.json');
+    const makeRequestStub = stub(request, 'makeRequest').resolves(
+      mockResponse as Response
+    );
+    const result = await generateContent(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams
+    );
+    expect(result.response.usageMetadata?.totalTokenCount).to.equal(1913);
+    expect(result.response.usageMetadata?.candidatesTokenCount).to.equal(76);
+    expect(result.response.usageMetadata?.promptTokensDetails?.[0].modality).to.equal('IMAGE');
+    expect(result.response.usageMetadata?.promptTokensDetails?.[0].tokenCount).to.equal(1806);
+    expect(result.response.usageMetadata?.candidatesTokensDetails?.[0].modality).to.equal('TEXT');
+    expect(result.response.usageMetadata?.candidatesTokensDetails?.[0].tokenCount).to.equal(76);
+    expect(makeRequestStub).to.be.calledWith(
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
+      match.any
+    );
+  });
   it('citations', async () => {
     const mockResponse = getMockResponse('unary-success-citations.json');
     const makeRequestStub = stub(request, 'makeRequest').resolves(

packages/vertexai/src/types/enums.ts

@@ -137,3 +137,16 @@ export enum FunctionCallingMode {
   // not passing any function declarations.
   NONE = 'NONE'
 }
+
+/**
+ * Type of content.
+ * @public
+ */
+export enum Modality {
+  MODALITY_UNSPECIFIED = 'MODALITY_UNSPECIFIED',
+  TEXT = 'TEXT',
+  IMAGE = 'IMAGE',
+  VIDEO = 'VIDEO',
+  AUDIO = 'AUDIO',
+  DOCUMENT = 'DOCUMENT'
+}

packages/vertexai/src/types/responses.ts

@@ -21,7 +21,8 @@ import {
   FinishReason,
   HarmCategory,
   HarmProbability,
-  HarmSeverity
+  HarmSeverity,
+  Modality
 } from './enums';
 
 /**
@@ -83,6 +84,17 @@ export interface UsageMetadata {
   promptTokenCount: number;
   candidatesTokenCount: number;
   totalTokenCount: number;
+  promptTokensDetails?: ModalityTokenCount[];
+  candidatesTokensDetails?: ModalityTokenCount[];
+}
+
+/**
+ * The number of tokens used by a given content type.
+ * @public
+ */
+export interface ModalityTokenCount {
+  modality: Modality;
+  tokenCount: number;
 }
 
 /**
@@ -213,4 +225,6 @@ export interface CountTokensResponse {
    * from the request.
    */
   totalBillableCharacters?: number;
+
+  promptTokensDetails?: ModalityTokenCount[];
 }

scripts/update_vertexai_responses.sh

@@ -17,7 +17,7 @@
 # This script replaces mock response files for Vertex AI unit tests with a fresh
 # clone of the shared repository of Vertex AI test data.
 
-RESPONSES_VERSION='v5.*' # The major version of mock responses to use
+RESPONSES_VERSION='v6.*' # The major version of mock responses to use
 REPO_NAME="vertexai-sdk-test-data"
 REPO_LINK="https://github.com/FirebaseExtended/$REPO_NAME.git"