feat(google-common), feat(core): Improve token counting (#8128)

Author: benjamincburns

langchain-core/src/messages/ai.ts

@@ -21,17 +21,41 @@ export type AIMessageFields = BaseMessageFields & {
   usage_metadata?: UsageMetadata;
 };
 
+export type ModalitiesTokenDetails = {
+  /**
+   * Text tokens.
+   * Does not need to be reported, but some models will do so.
+   */
+  text?: number;
+
+  /**
+   * Image (non-video) tokens.
+   */
+  image?: number;
+
+  /**
+   * Audio tokens.
+   */
+  audio?: number;
+
+  /**
+   * Video tokens.
+   */
+  video?: number;
+
+  /**
+   * Document tokens.
+   * e.g. PDF
+   */
+  document?: number;
+};
+
 /**
  * Breakdown of input token counts.
  *
  * Does not *need* to sum to full input token count. Does *not* need to have all keys.
  */
-export type InputTokenDetails = {
-  /**
-   * Audio input tokens.
-   */
-  audio?: number;
-
+export type InputTokenDetails = ModalitiesTokenDetails & {
   /**
    * Input tokens that were cached and there was a cache hit.
    *
@@ -53,12 +77,7 @@ export type InputTokenDetails = {
  *
  * Does *not* need to sum to full output token count. Does *not* need to have all keys.
  */
-export type OutputTokenDetails = {
-  /**
-   * Audio output tokens
-   */
-  audio?: number;
-
+export type OutputTokenDetails = ModalitiesTokenDetails & {
   /**
    * Reasoning output tokens.
    *

libs/langchain-google-common/src/types.ts

@@ -596,10 +596,39 @@ interface GeminiResponsePromptFeedback {
   safetyRatings: GeminiSafetyRating[];
 }
 
+export type ModalityEnum =
+  | "TEXT"
+  | "IMAGE"
+  | "VIDEO"
+  | "AUDIO"
+  | "DOCUMENT"
+  | string;
+
+export interface ModalityTokenCount {
+  modality: ModalityEnum;
+  tokenCount: number;
+}
+
+export interface GenerateContentResponseUsageMetadata {
+  promptTokenCount: number;
+  toolUsePromptTokenCount: number;
+  cachedContentTokenCount: number;
+  thoughtsTokenCount: number;
+  candidatesTokenCount: number;
+  totalTokenCount: number;
+
+  promptTokensDetails: ModalityTokenCount[];
+  toolUsePromptTokensDetails: ModalityTokenCount[];
+  cacheTokensDetails: ModalityTokenCount[];
+  candidatesTokensDetails: ModalityTokenCount[];
+
+  [key: string]: unknown;
+}
+
 export interface GenerateContentResponseData {
   candidates: GeminiResponseCandidate[];
   promptFeedback: GeminiResponsePromptFeedback;
-  usageMetadata: Record<string, unknown>;
+  usageMetadata: GenerateContentResponseUsageMetadata;
 }
 
 export type GoogleLLMModelFamily = null | "palm" | "gemini" | "gemma";

libs/langchain-google-common/src/utils/gemini.ts

@@ -19,6 +19,9 @@ import {
   parseBase64DataUrl,
   isDataContentBlock,
   convertToProviderContentBlock,
+  InputTokenDetails,
+  OutputTokenDetails,
+  ModalitiesTokenDetails,
 } from "@langchain/core/messages";
 import {
   ChatGeneration,
@@ -47,6 +50,7 @@ import type {
   GeminiLogprobsResult,
   GeminiLogprobsResultCandidate,
   GeminiLogprobsTopCandidate,
+  ModalityTokenCount,
 } from "../types.js";
 import { GoogleAISafetyError } from "./safety.js";
 import { MediaBlob } from "../experimental/utils/media_core.js";
@@ -855,6 +859,58 @@ export function getGeminiAPI(config?: GeminiAPIConfig): GoogleAIAPI {
     };
   }
 
+  function addModalityCounts(
+    modalityTokenCounts: ModalityTokenCount[],
+    details: InputTokenDetails | OutputTokenDetails
+  ): void {
+    modalityTokenCounts?.forEach((modalityTokenCount) => {
+      const { modality, tokenCount } = modalityTokenCount;
+      const modalityLc: keyof ModalitiesTokenDetails =
+        modality.toLowerCase() as keyof ModalitiesTokenDetails;
+      const currentCount = details[modalityLc] ?? 0;
+      // eslint-disable-next-line no-param-reassign
+      details[modalityLc] = currentCount + tokenCount;
+    });
+  }
+
+  function responseToUsageMetadata(
+    response: GoogleLLMResponse
+  ): UsageMetadata | undefined {
+    if ("usageMetadata" in response.data) {
+      const data: GenerateContentResponseData = response?.data;
+      const usageMetadata = data?.usageMetadata;
+
+      const input_tokens = usageMetadata.promptTokenCount ?? 0;
+      const candidatesTokenCount = usageMetadata.candidatesTokenCount ?? 0;
+      const thoughtsTokenCount = usageMetadata.thoughtsTokenCount ?? 0;
+      const output_tokens = candidatesTokenCount + thoughtsTokenCount;
+      const total_tokens =
+        usageMetadata.totalTokenCount ?? input_tokens + output_tokens;
+
+      const input_token_details: InputTokenDetails = {};
+      addModalityCounts(usageMetadata.promptTokensDetails, input_token_details);
+
+      const output_token_details: OutputTokenDetails = {};
+      addModalityCounts(
+        usageMetadata?.candidatesTokensDetails,
+        output_token_details
+      );
+      if (typeof usageMetadata?.thoughtsTokenCount === "number") {
+        output_token_details.reasoning = usageMetadata.thoughtsTokenCount;
+      }
+
+      const ret: UsageMetadata = {
+        input_tokens,
+        output_tokens,
+        total_tokens,
+        input_token_details,
+        output_token_details,
+      };
+      return ret;
+    }
+    return undefined;
+  }
+
   function responseToGenerationInfo(response: GoogleLLMResponse) {
     const data =
       // eslint-disable-next-line no-nested-ternary
@@ -890,11 +946,7 @@ export function getGeminiAPI(config?: GeminiAPIConfig): GoogleAIAPI {
     // Only add the usage_metadata on the last chunk
     // sent while streaming (see issue 8102).
     if (typeof finish_reason === "string") {
-      ret.usage_metadata = {
-        prompt_token_count: data.usageMetadata?.promptTokenCount,
-        candidates_token_count: data.usageMetadata?.candidatesTokenCount,
-        total_token_count: data.usageMetadata?.totalTokenCount,
-      };
+      ret.usage_metadata = responseToUsageMetadata(response);
     }
 
     return ret;
@@ -1115,15 +1167,7 @@ export function getGeminiAPI(config?: GeminiAPIConfig): GoogleAIAPI {
     const lastContent = gen.content[gen.content.length - 1];
 
     // Add usage metadata
-    let usageMetadata: UsageMetadata | undefined;
-    if ("usageMetadata" in response.data) {
-      usageMetadata = {
-        input_tokens: response.data.usageMetadata.promptTokenCount as number,
-        output_tokens: response.data.usageMetadata
-          .candidatesTokenCount as number,
-        total_tokens: response.data.usageMetadata.totalTokenCount as number,
-      };
-    }
+    const usage_metadata = responseToUsageMetadata(response);
 
     // Add thinking / reasoning
     // if (gen.reasoning && gen.reasoning.length > 0) {
@@ -1134,7 +1178,7 @@ export function getGeminiAPI(config?: GeminiAPIConfig): GoogleAIAPI {
     const message = new AIMessageChunk({
       content: combinedContent,
       additional_kwargs: kwargs,
-      usage_metadata: usageMetadata,
+      usage_metadata,
       tool_calls: combinedToolCalls.tool_calls,
       invalid_tool_calls: combinedToolCalls.invalid_tool_calls,
     });

libs/langchain-google-webauth/src/tests/chat_models.int.test.ts

@@ -39,6 +39,12 @@ import { ChatGoogle, ChatGoogleInput } from "../chat_models.js";
 import { BlobStoreAIStudioFile } from "../media.js";
 import MockedFunction = jest.MockedFunction;
 
+function propSum(o: Record<string, number>): number {
+  return Object.keys(o)
+    .map((key) => o[key])
+    .reduce((acc, val) => acc + val);
+}
+
 class WeatherTool extends StructuredTool {
   schema = z.object({
     locations: z
@@ -442,10 +448,16 @@ describe.each(testGeminiModelNames)(
       expect(text).toMatch(/(1 + 1 (equals|is|=) )?2.? ?/);
 
       expect(res).toHaveProperty("response_metadata");
-      expect(res.response_metadata).not.toHaveProperty("groundingMetadata");
-      expect(res.response_metadata).not.toHaveProperty("groundingSupport");
-
-      console.log(recorder);
+      const meta = res.response_metadata;
+      expect(meta).not.toHaveProperty("groundingMetadata");
+      expect(meta).not.toHaveProperty("groundingSupport");
+      expect(meta).toHaveProperty("usage_metadata");
+      const usage = meta.usage_metadata;
+
+      // Although LangChainJS doesn't require that the details sum to the
+      // available tokens, this should be the case for how we're doing Gemini.
+      expect(propSum(usage.input_token_details)).toEqual(usage.input_tokens);
+      expect(propSum(usage.output_token_details)).toEqual(usage.output_tokens);
     });
 
     test(`generate`, async () => {
@@ -883,6 +895,21 @@ describe.each(testGeminiModelNames)(
 
       expect(typeof response.content).toBe("string");
       expect((response.content as string).length).toBeGreaterThan(15);
+
+      expect(response).toHaveProperty("response_metadata");
+      const meta = response.response_metadata;
+      expect(meta).not.toHaveProperty("groundingMetadata");
+      expect(meta).not.toHaveProperty("groundingSupport");
+      expect(meta).toHaveProperty("usage_metadata");
+      const usage = meta.usage_metadata;
+
+      // Although LangChainJS doesn't require that the details sum to the
+      // available tokens, this should be the case for how we're doing Gemini.
+      expect(propSum(usage.input_token_details)).toEqual(usage.input_tokens);
+      expect(propSum(usage.output_token_details)).toEqual(usage.output_tokens);
+      expect(usage.input_token_details).toHaveProperty("audio");
+
+      console.log(response);
     });
 
     test("Supports GoogleSearchRetrievalTool", async () => {