fix: add reasoning tokens to usage stats for gemini models (#3339)

Author: MichaelDoyle

js/plugins/google-genai/src/common/types.ts

@@ -378,6 +378,8 @@ export declare interface UsageMetadata {
   totalTokenCount?: number;
   /** Optional. Number of tokens in the cached content. */
   cachedContentTokenCount?: number;
+  /** Optional. Number of tokens present in thoughts output. */
+  thoughtsTokenCount?: number;
 }
 
 export const TaskTypeSchema = z.enum([

js/plugins/google-genai/src/googleai/gemini.ts

@@ -621,6 +621,7 @@ export function defineModel(
             ...getBasicUsageStats(request.messages, candidateData),
             inputTokens: response.usageMetadata?.promptTokenCount,
             outputTokens: response.usageMetadata?.candidatesTokenCount,
+            thoughtsTokens: response.usageMetadata?.thoughtsTokenCount,
             totalTokens: response.usageMetadata?.totalTokenCount,
             cachedContentTokens:
               response.usageMetadata?.cachedContentTokenCount,

js/plugins/google-genai/src/vertexai/gemini.ts

@@ -643,6 +643,7 @@ export function defineModel(
             ...getBasicUsageStats(request.messages, candidateData),
             inputTokens: response.usageMetadata?.promptTokenCount,
             outputTokens: response.usageMetadata?.candidatesTokenCount,
+            thoughtsTokens: response.usageMetadata?.thoughtsTokenCount,
             totalTokens: response.usageMetadata?.totalTokenCount,
             cachedContentTokens:
               response.usageMetadata?.cachedContentTokenCount,

js/plugins/googleai/src/gemini.ts

@@ -36,6 +36,7 @@ import {
   type StartChatParams,
   type Tool,
   type ToolConfig,
+  type UsageMetadata,
 } from '@google/generative-ai';
 import {
   GENKIT_CLIENT_HEADER,
@@ -67,6 +68,13 @@ import { getApiKeyFromEnvVar } from './common';
 import { handleCacheIfNeeded } from './context-caching';
 import { extractCacheConfig } from './context-caching/utils';
 
+// Extra type guard to keep the compiler happy and avoid a cast to any. The
+// legacy Gemini SDK is no longer maintained, and doesn't have updated types.
+// However, the REST API returns the data we want.
+type ExtendedUsageMetadata = UsageMetadata & {
+  thoughtsTokenCount?: number;
+};
+
 /**
  * See https://ai.google.dev/gemini-api/docs/safety-settings#safety-filters.
  */
@@ -1385,16 +1393,18 @@ export function defineGoogleAIModel({
         const candidateData =
           candidates.map(fromJSONModeScopedGeminiCandidate) || [];
 
+        const usageMetadata = response.usageMetadata as ExtendedUsageMetadata;
+
         return {
           candidates: candidateData,
           custom: response,
           usage: {
             ...getBasicUsageStats(request.messages, candidateData),
-            inputTokens: response.usageMetadata?.promptTokenCount,
-            outputTokens: response.usageMetadata?.candidatesTokenCount,
-            totalTokens: response.usageMetadata?.totalTokenCount,
-            cachedContentTokens:
-              response.usageMetadata?.cachedContentTokenCount,
+            inputTokens: usageMetadata?.promptTokenCount,
+            outputTokens: usageMetadata?.candidatesTokenCount,
+            thoughtsTokens: usageMetadata?.thoughtsTokenCount,
+            totalTokens: usageMetadata?.totalTokenCount,
+            cachedContentTokens: usageMetadata?.cachedContentTokenCount,
           },
         };
       };

js/plugins/vertexai/src/gemini.ts

@@ -17,6 +17,7 @@
 import {
   FunctionCallingMode,
   FunctionDeclarationSchemaType,
+  UsageMetadata,
   type Content,
   type FunctionDeclaration,
   type Part as GeminiPart,
@@ -63,11 +64,17 @@ import {
 } from 'genkit/model/middleware';
 import { runInNewSpan } from 'genkit/tracing';
 import { GoogleAuth } from 'google-auth-library';
-
 import type { PluginOptions } from './common/types.js';
 import { handleCacheIfNeeded } from './context-caching/index.js';
 import { extractCacheConfig } from './context-caching/utils.js';
 
+// Extra type guard to keep the compiler happy and avoid a cast to any. The
+// legacy Gemini SDK is no longer maintained, and doesn't have updated types.
+// However, the REST API returns the data we want.
+type ExtendedUsageMetadata = UsageMetadata & {
+  thoughtsTokenCount?: number;
+};
+
 export const SafetySettingsSchema = z.object({
   category: z.enum([
     /** The harm category is unspecified. */
@@ -1276,16 +1283,18 @@ export function defineGeminiModel({
           fromGeminiCandidate(c, jsonMode)
         );
 
+        const usageMetadata = response.usageMetadata as ExtendedUsageMetadata;
+
         return {
           candidates: candidateData,
           custom: response,
           usage: {
             ...getBasicUsageStats(request.messages, candidateData),
-            inputTokens: response.usageMetadata?.promptTokenCount,
-            outputTokens: response.usageMetadata?.candidatesTokenCount,
-            totalTokens: response.usageMetadata?.totalTokenCount,
-            cachedContentTokens:
-              response.usageMetadata?.cachedContentTokenCount,
+            inputTokens: usageMetadata?.promptTokenCount,
+            outputTokens: usageMetadata?.candidatesTokenCount,
+            totalTokens: usageMetadata?.totalTokenCount,
+            thoughtsTokens: usageMetadata?.thoughtsTokenCount,
+            cachedContentTokens: usageMetadata?.cachedContentTokenCount,
           },
         };
       };