Add support for token-based usage metrics.

.changeset/chilled-tips-judge.md

@@ -0,0 +1,6 @@
+---
+'@firebase/vertexai': minor
+'firebase': minor
+---
+
+Added support for modality-based token count.

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,28 @@ 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 {
+    modality: Modality;
+    tokenCount: number;
+}
+
 // @public
 export interface ModelParams extends BaseParams {
     // (undocumented)
@@ -682,8 +706,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,33 @@ 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,40 @@ 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,22 @@ export enum FunctionCallingMode {
   // not passing any function declarations.
   NONE = 'NONE'
 }
+
+/**
+ * Content part modality.
+ * @public
+ */
+export enum Modality {
+  // Unspecified modality.
+  MODALITY_UNSPECIFIED = 'MODALITY_UNSPECIFIED',
+  // Plain text.
+  TEXT = 'TEXT',
+  // Image.
+  IMAGE = 'IMAGE',
+  // Video.
+  VIDEO = 'VIDEO',
+  // Audio.
+  AUDIO = 'AUDIO',
+  // Document, e.g. PDF.
+  DOCUMENT = 'DOCUMENT'
+}

packages/vertexai/src/types/responses.ts

@@ -21,7 +21,8 @@ import {
   FinishReason,
   HarmCategory,
   HarmProbability,
-  HarmSeverity
+  HarmSeverity,
+  Modality
 } from './enums';
 
 /**
@@ -83,6 +84,20 @@ export interface UsageMetadata {
   promptTokenCount: number;
   candidatesTokenCount: number;
   totalTokenCount: number;
+  promptTokensDetails?: ModalityTokenCount[];
+  candidatesTokensDetails?: ModalityTokenCount[];
+}
+
+/**
+ * Represents token counting info for a single modality.
+ *
+ * @public
+ */
+export interface ModalityTokenCount {
+  /** The modality associated with this token count. */
+  modality: Modality;
+  /** The number of tokens counted. */
+  tokenCount: number;
 }
 
 /**
@@ -213,4 +228,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"