Update usage metrics capture for gemini to include returned token cou…

…nts, and properly count different media types (images, video and audio files) (#332)

Gemini usage metrics now include returned token counts and correctly counts different media types (images, video and audio files). This was previously counting any media type as an image.

js/ai/src/model.ts

@@ -191,6 +191,10 @@ export const GenerationUsageSchema = z.object({
   outputCharacters: z.number().optional(),
   inputImages: z.number().optional(),
   outputImages: z.number().optional(),
+  inputVideos: z.number().optional(),
+  outputVideos: z.number().optional(),
+  inputAudioFiles: z.number().optional(),
+  outputAudioFiles: z.number().optional(),
   custom: z.record(z.number()).optional(),
 });
 export type GenerationUsage = z.infer<typeof GenerationUsageSchema>;
@@ -344,6 +348,8 @@ export function modelRef<
 type PartCounts = {
   characters: number;
   images: number;
+  videos: number;
+  audio: number;
 };
 
 /**
@@ -363,8 +369,12 @@ export function getBasicUsageStats(
   return {
     inputCharacters: inputCounts.characters,
     inputImages: inputCounts.images,
+    inputVideos: inputCounts.videos,
+    inputAudioFiles: inputCounts.audio,
     outputCharacters: outputCounts.characters,
     outputImages: outputCounts.images,
+    outputVideos: outputCounts.videos,
+    outputAudioFiles: outputCounts.audio,
   };
 }
 
@@ -373,10 +383,17 @@ function getPartCounts(parts: Part[]): PartCounts {
     (counts, part) => {
       return {
         characters: counts.characters + (part.text?.length || 0),
-        images: counts.images + (part.media ? 1 : 0),
+        images:
+          counts.images +
+          (part.media?.contentType?.startsWith('image') ? 1 : 0),
+        videos:
+          counts.videos +
+          (part.media?.contentType?.startsWith('video') ? 1 : 0),
+        audio:
+          counts.audio + (part.media?.contentType?.startsWith('audio') ? 1 : 0),
       };
     },
-    { characters: 0, images: 0 }
+    { characters: 0, images: 0, videos: 0, audio: 0 }
   );
 }
 

js/ai/src/telemetry.ts

@@ -28,6 +28,7 @@ import { GenerateOptions } from './generate.js';
 import {
   GenerateRequest,
   GenerateResponseData,
+  GenerationUsage,
   MediaPart,
   Part,
   ToolRequestPart,
@@ -77,6 +78,19 @@ const generateActionInputImages = new MetricCounter(
   }
 );
 
+const generateActionInputVideos = new MetricCounter(
+  _N('generate/input/videos'),
+  {
+    description: 'Counts input videos to a Genkit model.',
+    valueType: ValueType.INT,
+  }
+);
+
+const generateActionInputAudio = new MetricCounter(_N('generate/input/audio'), {
+  description: 'Counts input audio files to a Genkit model.',
+  valueType: ValueType.INT,
+});
+
 const generateActionOutputCharacters = new MetricCounter(
   _N('generate/output/characters'),
   {
@@ -101,6 +115,22 @@ const generateActionOutputImages = new MetricCounter(
   }
 );
 
+const generateActionOutputVideos = new MetricCounter(
+  _N('generate/output/videos'),
+  {
+    description: 'Count output videos from a Genkit model.',
+    valueType: ValueType.INT,
+  }
+);
+
+const generateActionOutputAudio = new MetricCounter(
+  _N('generate/output/audio'),
+  {
+    description: 'Count output audio files from a Genkit model.',
+    valueType: ValueType.INT,
+  }
+);
+
 type SharedDimensions = {
   modelName?: string;
   path?: string;
@@ -119,19 +149,12 @@ export function recordGenerateActionMetrics(
     err?: any;
   }
 ) {
-  doRecordGenerateActionMetrics(modelName, {
+  doRecordGenerateActionMetrics(modelName, opts.response?.usage || {}, {
     temperature: input.config?.temperature,
     topK: input.config?.topK,
     topP: input.config?.topP,
     maxOutputTokens: input.config?.maxOutputTokens,
     path: spanMetadataAls?.getStore()?.path,
-    inputTokens: opts.response?.usage?.inputTokens,
-    outputTokens: opts.response?.usage?.outputTokens,
-    totalTokens: opts.response?.usage?.totalTokens,
-    inputCharacters: opts.response?.usage?.inputCharacters,
-    outputCharacters: opts.response?.usage?.outputCharacters,
-    inputImages: opts.response?.usage?.inputImages,
-    outputImages: opts.response?.usage?.outputImages,
     latencyMs: opts.response?.latencyMs,
     err: opts.err,
     source: 'ts',
@@ -290,20 +313,13 @@ function toPartLogToolResponse(part: ToolResponsePart): string {
  */
 function doRecordGenerateActionMetrics(
   modelName: string,
+  usage: GenerationUsage,
   dimensions: {
     path?: string;
     temperature?: number;
     maxOutputTokens?: number;
     topK?: number;
     topP?: number;
-    inputTokens?: number;
-    outputTokens?: number;
-    totalTokens?: number;
-    inputCharacters?: number;
-    outputCharacters?: number;
-    totalCharacters?: number;
-    inputImages?: number;
-    outputImages?: number;
     latencyMs?: number;
     err?: any;
     source?: string;
@@ -329,12 +345,16 @@ function doRecordGenerateActionMetrics(
   generateActionLatencies.record(dimensions.latencyMs, shared);
 
   // inputs
-  generateActionInputTokens.add(dimensions.inputTokens, shared);
-  generateActionInputCharacters.add(dimensions.inputCharacters, shared);
-  generateActionInputImages.add(dimensions.inputImages, shared);
+  generateActionInputTokens.add(usage.inputTokens, shared);
+  generateActionInputCharacters.add(usage.inputCharacters, shared);
+  generateActionInputImages.add(usage.inputImages, shared);
+  generateActionInputVideos.add(usage.inputVideos, shared);
+  generateActionInputAudio.add(usage.inputAudioFiles, shared);
 
   // outputs
-  generateActionOutputTokens.add(dimensions.outputTokens, shared);
-  generateActionOutputCharacters.add(dimensions.outputCharacters, shared);
-  generateActionOutputImages.add(dimensions.outputImages, shared);
+  generateActionOutputTokens.add(usage.outputTokens, shared);
+  generateActionOutputCharacters.add(usage.outputCharacters, shared);
+  generateActionOutputImages.add(usage.outputImages, shared);
+  generateActionOutputVideos.add(usage.outputVideos, shared);
+  generateActionOutputAudio.add(usage.outputAudioFiles, shared);
 }

js/plugins/googleai/src/gemini.ts

@@ -500,7 +500,12 @@ export function googleAIModel(
         return {
           candidates: responseCandidates,
           custom: result.response,
-          usage: getBasicUsageStats(request.messages, responseCandidates),
+          usage: {
+            ...getBasicUsageStats(request.messages, responseCandidates),
+            inputTokens: result.response.usageMetadata?.promptTokenCount,
+            outputTokens: result.response.usageMetadata?.candidatesTokenCount,
+            totalTokens: result.response.usageMetadata?.totalTokenCount,
+          },
         };
       }
     }

js/plugins/vertexai/src/gemini.ts

@@ -513,7 +513,12 @@ export function geminiModel(name: string, vertex: VertexAI): ModelAction {
         return {
           candidates: responseCandidates,
           custom: result.response,
-          usage: getBasicUsageStats(request.messages, responseCandidates),
+          usage: {
+            ...getBasicUsageStats(request.messages, responseCandidates),
+            inputTokens: result.response.usageMetadata?.promptTokenCount,
+            outputTokens: result.response.usageMetadata?.candidatesTokenCount,
+            totalTokens: result.response.usageMetadata?.totalTokenCount,
+          },
         };
       }
     }