From b5aa892fafa7718171dfecb18ecd01d88e878918 Mon Sep 17 00:00:00 2001 From: Bryan Atkinson Date: Mon, 10 Jun 2024 13:37:04 -0400 Subject: [PATCH] Update usage metrics capture for gemini to include returned token counts, 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 | 21 +++++++++- js/ai/src/telemetry.ts | 64 ++++++++++++++++++++----------- js/plugins/googleai/src/gemini.ts | 7 +++- js/plugins/vertexai/src/gemini.ts | 7 +++- 4 files changed, 73 insertions(+), 26 deletions(-) diff --git a/js/ai/src/model.ts b/js/ai/src/model.ts index 10036ea62..1d8acd57c 100644 --- a/js/ai/src/model.ts +++ b/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; @@ -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 } ); } diff --git a/js/ai/src/telemetry.ts b/js/ai/src/telemetry.ts index 59d5bbf87..b8e7ae6f3 100644 --- a/js/ai/src/telemetry.ts +++ b/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); } diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts index 8c12710e9..0a6802a1b 100644 --- a/js/plugins/googleai/src/gemini.ts +++ b/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, + }, }; } } diff --git a/js/plugins/vertexai/src/gemini.ts b/js/plugins/vertexai/src/gemini.ts index 0dada89b9..4bd1c3f08 100644 --- a/js/plugins/vertexai/src/gemini.ts +++ b/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, + }, }; } }