|
| 1 | +/* |
| 2 | + * Copyright 2026 New Relic Corporation. All rights reserved. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +const LlmChatCompletionSummary = require('../chat-summary') |
| 7 | +const { getUsageTokens } = require('./utils') |
| 8 | + |
| 9 | +/** |
| 10 | + * Encapsulates a Google Gen AI LlmChatCompletionSummary. |
| 11 | + */ |
| 12 | +class GoogleGenAiLlmChatCompletionSummary extends LlmChatCompletionSummary { |
| 13 | + constructor({ agent, segment, transaction, request, response, error }) { |
| 14 | + super({ agent, |
| 15 | + segment, |
| 16 | + transaction, |
| 17 | + responseModel: response?.modelVersion, |
| 18 | + requestModel: request?.model, |
| 19 | + finishReason: response?.candidates?.[0]?.finishReason, |
| 20 | + maxTokens: request.config?.maxOutputTokens, |
| 21 | + temperature: request.config?.temperature, |
| 22 | + vendor: 'gemini', |
| 23 | + error }) |
| 24 | + |
| 25 | + let requestMessagesLength = 0 |
| 26 | + if (Array.isArray(request?.contents)) { |
| 27 | + requestMessagesLength = request.contents.length |
| 28 | + } else if (typeof request?.contents === 'string') { |
| 29 | + requestMessagesLength = 1 |
| 30 | + } |
| 31 | + this['response.number_of_messages'] = requestMessagesLength + (response?.candidates?.length || 0) |
| 32 | + |
| 33 | + this.setTokens(agent, request, response) |
| 34 | + } |
| 35 | + |
| 36 | + setTokens(agent, request, response) { |
| 37 | + const tokenCB = agent.llm?.tokenCountCallback |
| 38 | + |
| 39 | + // Prefer callback for prompt and completion tokens; if unavailable, fall back to response data. |
| 40 | + if (tokenCB) { |
| 41 | + const promptContent = typeof request?.contents === 'string' |
| 42 | + ? request?.contents |
| 43 | + : request?.contents?.join(' ') |
| 44 | + |
| 45 | + const responseContent = response?.candidates?.[0]?.content?.parts |
| 46 | + const completionContent = responseContent?.map((content) => content.text).join(' ') |
| 47 | + |
| 48 | + if (promptContent && completionContent) { |
| 49 | + this.setTokenUsageFromCallback( |
| 50 | + { |
| 51 | + tokenCB, |
| 52 | + reqModel: request.model, |
| 53 | + resModel: this['response.model'], |
| 54 | + promptContent, |
| 55 | + completionContent |
| 56 | + } |
| 57 | + ) |
| 58 | + } |
| 59 | + return |
| 60 | + } |
| 61 | + |
| 62 | + const tokens = getUsageTokens(response) |
| 63 | + this.setTokensInResponse(tokens) |
| 64 | + } |
| 65 | +} |
| 66 | + |
| 67 | +module.exports = GoogleGenAiLlmChatCompletionSummary |
0 commit comments