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

docs-devsite/_toc.yaml

@@ -540,6 +540,8 @@ toc:
     path: /docs/reference/js/vertexai.inlinedatapart.md
   - title: IntegerSchema
     path: /docs/reference/js/vertexai.integerschema.md
+  - title: ModalityTokenCount
+    path: /docs/reference/js/vertexai.modalitytokencount.md
   - title: ModelParams
     path: /docs/reference/js/vertexai.modelparams.md
   - title: NumberSchema

docs-devsite/vertexai.counttokensresponse.md

@@ -22,9 +22,20 @@ export interface CountTokensResponse
 
 |  Property | Type | Description |
 |  --- | --- | --- |
+|  [promptTokensDetails](./vertexai.counttokensresponse.md#counttokensresponseprompttokensdetails) | [ModalityTokenCount](./vertexai.modalitytokencount.md#modalitytokencount_interface)<!-- -->\[\] | The breakdown, by modality, of how many tokens are consumed by the prompt. |
 |  [totalBillableCharacters](./vertexai.counttokensresponse.md#counttokensresponsetotalbillablecharacters) | number | The total number of billable characters counted across all instances from the request. |
 |  [totalTokens](./vertexai.counttokensresponse.md#counttokensresponsetotaltokens) | number | The total number of tokens counted across all instances from the request. |
 
+## CountTokensResponse.promptTokensDetails
+
+The breakdown, by modality, of how many tokens are consumed by the prompt.
+
+<b>Signature:</b>
+
+```typescript
+promptTokensDetails?: ModalityTokenCount[];
+```
+
 ## CountTokensResponse.totalBillableCharacters
 
 The total number of billable characters counted across all instances from the request.

docs-devsite/vertexai.md

@@ -48,6 +48,7 @@ The Vertex AI in Firebase Web SDK.
 |  [HarmCategory](./vertexai.md#harmcategory) | Harm categories that would cause prompts or candidates to be blocked. |
 |  [HarmProbability](./vertexai.md#harmprobability) | Probability that a prompt or candidate matches a harm category. |
 |  [HarmSeverity](./vertexai.md#harmseverity) | Harm severity levels. |
+|  [Modality](./vertexai.md#modality) | Content part modality. |
 |  [SchemaType](./vertexai.md#schematype) | Contains the list of OpenAPI data types as defined by the [OpenAPI specification](https://swagger.io/docs/specification/data-models/data-types/) |
 |  [VertexAIErrorCode](./vertexai.md#vertexaierrorcode) | Standardized error codes that <code>[VertexAIError](./vertexai.vertexaierror.md#vertexaierror_class)</code> can have. |
 
@@ -84,6 +85,7 @@ The Vertex AI in Firebase Web SDK.
 |  [GroundingAttribution](./vertexai.groundingattribution.md#groundingattribution_interface) |  |
 |  [GroundingMetadata](./vertexai.groundingmetadata.md#groundingmetadata_interface) | Metadata returned to client when grounding is enabled. |
 |  [InlineDataPart](./vertexai.inlinedatapart.md#inlinedatapart_interface) | Content part interface if the part represents an image. |
+|  [ModalityTokenCount](./vertexai.modalitytokencount.md#modalitytokencount_interface) | Represents token counting info for a single modality. |
 |  [ModelParams](./vertexai.modelparams.md#modelparams_interface) | Params passed to <code>[getGenerativeModel()](./vertexai.md#getgenerativemodel_e3037c9)</code>. |
 |  [ObjectSchemaInterface](./vertexai.objectschemainterface.md#objectschemainterface_interface) | Interface for <code>[ObjectSchema](./vertexai.objectschema.md#objectschema_class)</code> class. |
 |  [PromptFeedback](./vertexai.promptfeedback.md#promptfeedback_interface) | If the prompt was blocked, this will be populated with <code>blockReason</code> and the relevant <code>safetyRatings</code>. |
@@ -363,6 +365,27 @@ export declare enum HarmSeverity
 |  HARM\_SEVERITY\_MEDIUM | <code>&quot;HARM_SEVERITY_MEDIUM&quot;</code> |  |
 |  HARM\_SEVERITY\_NEGLIGIBLE | <code>&quot;HARM_SEVERITY_NEGLIGIBLE&quot;</code> |  |
 
+## Modality
+
+Content part modality.
+
+<b>Signature:</b>
+
+```typescript
+export declare enum Modality 
+```
+
+## Enumeration Members
+
+|  Member | Value | Description |
+|  --- | --- | --- |
+|  AUDIO | <code>&quot;AUDIO&quot;</code> | Audio. |
+|  DOCUMENT | <code>&quot;DOCUMENT&quot;</code> | Document, e.g. PDF. |
+|  IMAGE | <code>&quot;IMAGE&quot;</code> | Image. |
+|  MODALITY\_UNSPECIFIED | <code>&quot;MODALITY_UNSPECIFIED&quot;</code> | Unspecified modality. |
+|  TEXT | <code>&quot;TEXT&quot;</code> | Plain text. |
+|  VIDEO | <code>&quot;VIDEO&quot;</code> | Video. |
+
 ## SchemaType
 
 Contains the list of OpenAPI data types as defined by the [OpenAPI specification](https://swagger.io/docs/specification/data-models/data-types/)

docs-devsite/vertexai.usagemetadata.md

@@ -23,7 +23,9 @@ export interface UsageMetadata
 |  Property | Type | Description |
 |  --- | --- | --- |
 |  [candidatesTokenCount](./vertexai.usagemetadata.md#usagemetadatacandidatestokencount) | number |  |
+|  [candidatesTokensDetails](./vertexai.usagemetadata.md#usagemetadatacandidatestokensdetails) | [ModalityTokenCount](./vertexai.modalitytokencount.md#modalitytokencount_interface)<!-- -->\[\] |  |
 |  [promptTokenCount](./vertexai.usagemetadata.md#usagemetadataprompttokencount) | number |  |
+|  [promptTokensDetails](./vertexai.usagemetadata.md#usagemetadataprompttokensdetails) | [ModalityTokenCount](./vertexai.modalitytokencount.md#modalitytokencount_interface)<!-- -->\[\] |  |
 |  [totalTokenCount](./vertexai.usagemetadata.md#usagemetadatatotaltokencount) | number |  |
 
 ## UsageMetadata.candidatesTokenCount
@@ -34,6 +36,14 @@ export interface UsageMetadata
 candidatesTokenCount: number;
 ```
 
+## UsageMetadata.candidatesTokensDetails
+
+<b>Signature:</b>
+
+```typescript
+candidatesTokensDetails?: ModalityTokenCount[];
+```
+
 ## UsageMetadata.promptTokenCount
 
 <b>Signature:</b>
@@ -42,6 +52,14 @@ candidatesTokenCount: number;
 promptTokenCount: number;
 ```
 
+## UsageMetadata.promptTokensDetails
+
+<b>Signature:</b>
+
+```typescript
+promptTokensDetails?: ModalityTokenCount[];
+```
+
 ## UsageMetadata.totalTokenCount
 
 <b>Signature:</b>

packages/rules-unit-testing/api-extractor.json

@@ -1,5 +1,5 @@
 {
     "extends": "../../config/api-extractor.json",
-    // Point it to your entry point d.ts file. 
-    "mainEntryPointFilePath": "<projectFolder>/dist/rules-unit-testing/index.d.ts"
-}
+    // Point it to your entry point d.ts file.
+    "mainEntryPointFilePath": "<projectFolder>/dist/index.d.ts"
+}

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,34 @@ 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 (for example, 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,8 @@ export interface CountTokensResponse {
    * from the request.
    */
   totalBillableCharacters?: number;
+  /**
+   * The breakdown, by modality, of how many tokens are consumed by the prompt.
+   */
+  promptTokensDetails?: ModalityTokenCount[];
 }