newrelic 13.6.2 → 13.6.3

package/lib/llm-events/aws-bedrock/stream-handler.js

@@ -187,6 +187,13 @@ class StreamHandler {
     this.response.response.headers = {
       'x-amzn-requestid': this.passThroughParams.response.response.headers['x-amzn-requestid']
     }
+
+    if (parsedEvent['amazon-bedrock-invocationMetrics']) {
+      const invocationMetrics = parsedEvent['amazon-bedrock-invocationMetrics']
+      this.response.response.headers['x-amzn-bedrock-input-token-count'] = invocationMetrics?.inputTokenCount
+      this.response.response.headers['x-amzn-bedrock-output-token-count'] = invocationMetrics?.outputTokenCount
+    }
+
     delete parsedEvent['amazon-bedrock-invocationMetrics']
   }
 }
@@ -213,8 +220,7 @@ async function * handleClaude() {
 async function * handleClaude3() {
   let currentBody = {}
   let stopReason
-  const completions = []
-  const usage = {}
+  let response = ''
 
   try {
     for await (const event of this.stream) {
@@ -223,25 +229,14 @@ async function * handleClaude3() {
       this.updateHeaders(parsed)
       currentBody = parsed
       if (parsed.type === 'content_block_delta') {
-        completions.push(parsed.delta.text)
+        response += parsed.delta.text
       } else if (parsed.type === 'message_delta') {
         stopReason = parsed.delta.stop_reason
       }
-
-      if (parsed?.message?.usage?.input_tokens) {
-        usage.input_tokens = parsed.message.usage.input_tokens
-      }
-
-      if (parsed?.message?.usage?.output_tokens) {
-        usage.output_tokens = parsed.message.usage.output_tokens
-      } else if (parsed?.usage?.output_tokens) {
-        usage.output_tokens = parsed.usage.output_tokens
-      }
     }
   } finally {
-    currentBody.completions = completions
+    currentBody.completions = response
     currentBody.stop_reason = stopReason
-    currentBody.usage = usage
     this.response.output.body = currentBody
     this.finish()
   }
@@ -286,7 +281,6 @@ async function * handleCohereEmbed() {
 async function * handleLlama() {
   let currentBody = {}
   let generation = ''
-  const usage = {}
 
   try {
     for await (const event of this.stream) {
@@ -295,35 +289,34 @@ async function * handleLlama() {
       this.updateHeaders(parsed)
       currentBody = parsed
       generation += parsed.generation
-
-      if (parsed?.generation_token_count) {
-        usage.output_tokens = parsed.generation_token_count
-      }
-
-      if (parsed?.prompt_token_count) {
-        usage.input_tokens = parsed.prompt_token_count
-      }
     }
   } finally {
     currentBody.generation = generation
-    currentBody.usage = usage
     this.response.output.body = currentBody
     this.finish()
   }
 }
 
+/**
+ * Yields every chunk and builds up the response in a string
+ * Re-assigns the response to `outputText` and makes body.results an array
+ * to match the structure of non-streamed Titan response.
+ */
 async function * handleTitan() {
-  const body = this.response.output.body
-  body.results = []
+  let currentBody = {}
+  let response = ''
 
   try {
     for await (const event of this.stream) {
       yield event // Pass it up to the real consumer of the stream.
       const parsed = this.parseEvent(event)
       this.updateHeaders(parsed)
-      body.results.push(parsed)
+      response += parsed.outputText
+      currentBody = parsed
     }
   } finally {
+    currentBody.outputText = response
+    this.response.output.body.results = [currentBody]
     this.finish()
   }
 }

package/lib/llm-events/aws-bedrock/utils.js

@@ -5,8 +5,6 @@
 
 'use strict'
 
-const { setTokensInResponse } = require('../utils')
-
 /**
  *
  * @param {object[]} chunks - The "chunks" that make up a single conceptual message. In a multi-modal scenario, a single message
@@ -69,63 +67,7 @@ function stringifyConverseChunkedMessage(chunks) {
   return stringifiedChunks.join('\n\n')
 }
 
-/**
- * Set the prompt, completion and total token counts on the Llm Event if input and output metadata
- * exists in the response
- * @param {object} response - The response object.
- * @param {object} context - The Llm Event
- */
-function setUsageTokens(response, context) {
-  // input and output token counts must available in order to add all usage attributes to response
-  // if total tokens is not available, we can manually add it up (from input and output token count)
-  if (tokenUsageAttributesExist(response)) {
-    const promptTokens =
-      Number(response?.usage?.input_tokens ||
-       response?.usage?.inputTokens ||
-        response?.parsedBody?.results?.[0]?.inputTextTokenCount ||
-         response?.parsedBody?.usage?.input_tokens)
-    const completionTokens =
-        Number(response?.usage?.output_tokens ||
-         response?.usage?.outputTokens ||
-          response?.parsedBody?.results?.[1]?.totalOutputTextTokenCount ||
-           response?.parsedBody?.usage.output_tokens)
-    const totalTokens = Number(response?.usage?.total_tokens || response?.usage?.totalTokens)
-
-    setTokensInResponse(context, { promptTokens, completionTokens, totalTokens })
-    return
-  }
-
-  if (tokenUsageHeadersExist(response)) {
-    const promptTokens = Number(response?.headers['x-amzn-bedrock-input-token-count'])
-    const completionTokens = Number(response?.headers['x-amzn-bedrock-output-token-count'])
-    const totalTokens = Number(response?.headers['x-amzn-bedrock-total-token-count'])
-
-    setTokensInResponse(context, { promptTokens, completionTokens, totalTokens })
-  }
-}
-
-function tokenUsageAttributesExist(response) {
-  const tokensA = response?.usage?.input_tokens && response?.usage?.output_tokens
-  const tokensB = response?.usage?.inputTokens && response?.usage?.outputTokens
-
-  // titan stream puts tokens in here
-  const tokensC = response?.parsedBody?.results?.[0]?.inputTextTokenCount && response?.parsedBody?.results?.[1]?.totalOutputTextTokenCount
-
-  // for llama3 and claude 3 streaming
-  const tokensD = response?.parsedBody?.usage?.input_tokens && response?.parsedBody?.usage?.output_tokens
-
-  return tokensA !== undefined || tokensB !== undefined || tokensC !== undefined || tokensD !== undefined
-}
-
-function tokenUsageHeadersExist(response) {
-  const tokens = response?.headers['x-amzn-bedrock-input-token-count'] && response?.headers['x-amzn-bedrock-output-token-count']
-  return tokens !== undefined
-}
-
 module.exports = {
   stringifyClaudeChunkedMessage,
-  stringifyConverseChunkedMessage,
-  tokenUsageHeadersExist,
-  tokenUsageAttributesExist,
-  setUsageTokens,
+  stringifyConverseChunkedMessage
 }

package/lib/llm-events/event.js

@@ -18,6 +18,114 @@ class BaseLlmEvent {
       }
     }
   }
+
+  /**
+   * Determines if the provided token count is valid.
+   * A valid token count is greater than 0 and not null.
+   * @param {number} tokenCount The token count obtained from the token callback
+   * @returns {boolean} Whether the token count is valid
+   */
+  validTokenCount(tokenCount) {
+    return tokenCount !== null && tokenCount > 0
+  }
+
+  /**
+   * Calculates the total token count from the prompt tokens and completion tokens
+   * set in the event.
+   * @returns {number} The total token count
+   */
+  getTotalTokenCount() {
+    return Number(this['response.usage.prompt_tokens']) + Number(this['response.usage.completion_tokens'])
+  }
+
+  setTokensOnEmbeddingMessage(totalTokens) {
+    if (this.validTokenCount(totalTokens)) {
+      this['response.usage.total_tokens'] = totalTokens
+    }
+  }
+
+  /**
+   * Sets the provided tokens counts on the LLM event.
+   * Checks if promptToken and completionToken are greater than zero before setting.
+   * This is because the spec states that token counts should only be set if both
+   * are present.
+   * @param {object} params to the function
+   * @param {object} params.promptTokens value of prompt token count
+   * @param {object} params.completionTokens value of completion(s) token count
+   * @param {object} params.totalTokens value of prompt + completion(s) token count
+   */
+  setTokensInResponse({ promptTokens, completionTokens, totalTokens }) {
+    if (this.validTokenCount(promptTokens) && this.validTokenCount(completionTokens)) {
+      this['response.usage.prompt_tokens'] = promptTokens
+      this['response.usage.completion_tokens'] = completionTokens
+      this['response.usage.total_tokens'] = totalTokens || this.getTotalTokenCount()
+    }
+  }
+
+  /**
+   * Sets `token_count` to 0 on the LlmChatCompletionMessage if both prompt and completion tokens are greater than zero.
+   * This is because the spec states that if token counts are set, then we should set token_count to 0 to indicate
+   * that the token calculation does not have to occur in the ingest pipeline.
+   * @param {object} params to the function
+   * @param {object} params.promptTokens value of prompt token count
+   * @param {object} params.completionTokens value of completion(s) token count
+   */
+  setTokenInCompletionMessage({ promptTokens, completionTokens }) {
+    if (this.validTokenCount(promptTokens) && this.validTokenCount(completionTokens)) {
+      this.token_count = 0
+    }
+  }
+
+  /**
+   * Calculates prompt and completion token counts using the provided callback and models.
+   * If both counts are valid, sets this.token_count to 0.
+   *
+   * @param {object} options - The params object.
+   * @param {Function} options.tokenCB - The token counting callback function.
+   * @param {string} options.reqModel - The model used for the prompt.
+   * @param {string} options.resModel - The model used for the completion.
+   * @param {string} options.promptContent - The prompt content to count tokens for.
+   * @param {string} options.completionContent - The completion content to count tokens for.
+   * @returns {void}
+   */
+  setTokenFromCallback({ tokenCB, reqModel, resModel, promptContent, completionContent }) {
+    const promptToken = this.calculateCallbackTokens(tokenCB, reqModel, promptContent)
+    const completionToken = this.calculateCallbackTokens(tokenCB, resModel, completionContent)
+
+    this.setTokenInCompletionMessage({ promptTokens: promptToken, completionTokens: completionToken })
+  }
+
+  /**
+   * Calculates prompt and completion token counts using the provided callback and models.
+   * If both counts are valid, sets token prompt, completion and total counts on the event.
+   *
+   * @param {object} options - The params object.
+   * @param {Function} options.tokenCB - The token counting callback function.
+   * @param {string} options.reqModel - The model used for the prompt.
+   * @param {string} options.resModel - The model used for the completion.
+   * @param {string} options.promptContent - The prompt content to count tokens for.
+   * @param {string} options.completionContent - The completion content to count tokens for.
+   * @returns {void}
+   */
+  setTokenUsageFromCallback({ tokenCB, reqModel, resModel, promptContent, completionContent }) {
+    const promptTokens = this.calculateCallbackTokens(tokenCB, reqModel, promptContent)
+    const completionTokens = this.calculateCallbackTokens(tokenCB, resModel, completionContent)
+    this.setTokensInResponse({ promptTokens, completionTokens, totalTokens: promptTokens + completionTokens })
+  }
+
+  /**
+   * Calculate the token counts using the provided callback.
+   * @param {Function} tokenCB - The token count callback function.
+   * @param {string} model - The model.
+   * @param {string} content - The content to calculate tokens for, such as prompt or completion response.
+   * @returns {number|undefined} - The calculated token count or undefined if callback is not a function.
+   */
+  calculateCallbackTokens(tokenCB, model, content) {
+    if (typeof tokenCB === 'function') {
+      return tokenCB(model, content)
+    }
+    return undefined
+  }
 }
 
 module.exports = BaseLlmEvent

package/lib/llm-events/google-genai/chat-completion-message.js

@@ -6,8 +6,6 @@
 'use strict'
 const LlmEvent = require('./event')
 const { makeId } = require('../../util/hashes')
-const { tokenUsageAttributesExist } = require('./utils')
-const { setTokenFromCallback } = require('../utils')
 
 module.exports = class LlmChatCompletionMessage extends LlmEvent {
   constructor({
@@ -50,9 +48,8 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
       const completionContent = responseContent?.map((content) => content.text).join(' ')
 
       if (promptContent && completionContent) {
-        setTokenFromCallback(
+        this.setTokenFromCallback(
           {
-            context: this,
             tokenCB,
             reqModel: request.model,
             resModel: this['response.model'],
@@ -64,11 +61,7 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
       return
     }
 
-    // If no token count callback is available, we need to check the response object
-    // for usage information and set token_count to 0 if all usage attributes are present.
-    // Response headers won't have token usage information
-    if (tokenUsageAttributesExist(response)) {
-      this.token_count = 0
-    }
+    const tokens = this.getUsageTokens(response)
+    this.setTokenInCompletionMessage(tokens)
   }
 }

package/lib/llm-events/google-genai/chat-completion-summary.js

@@ -5,8 +5,6 @@
 
 'use strict'
 const LlmEvent = require('./event')
-const { setUsageTokens } = require('./utils')
-const { setTokenUsageFromCallback } = require('../utils')
 
 module.exports = class LlmChatCompletionSummary extends LlmEvent {
   constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -39,9 +37,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
       const completionContent = responseContent?.map((content) => content.text).join(' ')
 
       if (promptContent && completionContent) {
-        setTokenUsageFromCallback(
+        this.setTokenUsageFromCallback(
           {
-            context: this,
             tokenCB,
             reqModel: request.model,
             resModel: this['response.model'],
@@ -52,6 +49,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
       }
       return
     }
-    setUsageTokens(response, this)
+
+    const tokens = this.getUsageTokens(response)
+    this.setTokensInResponse(tokens)
   }
 }

package/lib/llm-events/google-genai/embedding.js

@@ -6,7 +6,6 @@
 'use strict'
 
 const LlmEvent = require('./event')
-const { validCallbackTokenCount, calculateCallbackTokens } = require('../utils')
 
 class LlmEmbedding extends LlmEvent {
   constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -16,39 +15,6 @@ class LlmEmbedding extends LlmEvent {
     if (agent.config.ai_monitoring.record_content.enabled === true) {
       this.input = request.contents?.toString()
     }
-
-    this.setTotalTokens(agent, request, response)
-  }
-
-  setTotalTokens(agent, request, response) {
-    const tokenCB = agent.llm?.tokenCountCallback
-
-    // For embedding events, only total token count is relevant.
-    // Prefer callback for total tokens; if unavailable, fall back to response data.
-    if (tokenCB) {
-      const content = request.contents?.toString()
-
-      if (content === undefined) {
-        return
-      }
-
-      const totalTokenCount = calculateCallbackTokens(tokenCB, this['request.model'], content)
-      const hasValidCallbackCounts = validCallbackTokenCount(totalTokenCount)
-
-      if (hasValidCallbackCounts) {
-        this['response.usage.total_tokens'] = Number(totalTokenCount)
-      }
-      return
-    }
-
-    const totalTokens = this.getTotalTokens(response)
-    if (totalTokens) {
-      this['response.usage.total_tokens'] = Number(totalTokens)
-    }
-  }
-
-  getTotalTokens(response) {
-    return response?.usageMetadata?.totalTokenCount
   }
 }
 

package/lib/llm-events/google-genai/event.js

@@ -26,4 +26,11 @@ module.exports = class LlmEvent extends BaseEvent {
       this.duration = segment?.getDurationInMillis()
     }
   }
+
+  getUsageTokens(response) {
+    const promptTokens = Number(response?.usageMetadata?.promptTokenCount)
+    const completionTokens = Number(response?.usageMetadata?.candidatesTokenCount)
+    const totalTokens = Number(response?.usageMetadata?.totalTokenCount)
+    return { promptTokens, completionTokens, totalTokens }
+  }
 }

package/lib/llm-events/openai/chat-completion-message.js

@@ -5,8 +5,6 @@
 
 'use strict'
 const LlmEvent = require('./event')
-const { tokenUsageAttributesExist } = require('./utils')
-const { setTokenFromCallback } = require('../utils')
 
 module.exports = class LlmChatCompletionMessage extends LlmEvent {
   constructor({
@@ -58,9 +56,8 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
         : response?.choices?.map((resContent) => resContent.message.content).join(' ')
 
       if (promptContent && completionContent) {
-        setTokenFromCallback(
+        this.setTokenFromCallback(
           {
-            context: this,
             tokenCB,
             reqModel: request.model,
             resModel: this['response.model'],
@@ -72,10 +69,7 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
       return
     }
 
-    // If no token count callback is available, we need to check the response object
-    // for usage information and set token_count to 0 if all usage attributes are present.
-    if (tokenUsageAttributesExist(response)) {
-      this.token_count = 0
-    }
+    const tokens = this.getUsageTokens(response)
+    this.setTokenInCompletionMessage(tokens)
   }
 }

package/lib/llm-events/openai/chat-completion-summary.js

@@ -5,8 +5,6 @@
 
 'use strict'
 const LlmEvent = require('./event')
-const { setUsageTokens } = require('./utils')
-const { setTokenUsageFromCallback } = require('../utils')
 
 module.exports = class LlmChatCompletionSummary extends LlmEvent {
   constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -46,9 +44,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
         : response?.choices?.map((resContent) => resContent.message.content).join(' ')
 
       if (promptContent && completionContent) {
-        setTokenUsageFromCallback(
+        this.setTokenUsageFromCallback(
           {
-            context: this,
             tokenCB,
             reqModel: request.model,
             resModel: this['response.model'],
@@ -59,6 +56,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
       }
       return
     }
-    setUsageTokens(response, this)
+
+    const tokens = this.getUsageTokens(response)
+    this.setTokensInResponse(tokens)
   }
 }

package/lib/llm-events/openai/embedding.js

@@ -5,7 +5,6 @@
 
 'use strict'
 const LlmEvent = require('./event')
-const { validCallbackTokenCount, calculateCallbackTokens } = require('../utils')
 
 module.exports = class LlmEmbedding extends LlmEvent {
   constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -31,22 +30,12 @@ module.exports = class LlmEmbedding extends LlmEvent {
         return
       }
 
-      const totalTokenCount = calculateCallbackTokens(tokenCB, this['request.model'], content)
-      const hasValidCallbackCounts = validCallbackTokenCount(totalTokenCount)
-
-      if (hasValidCallbackCounts) {
-        this['response.usage.total_tokens'] = Number(totalTokenCount)
-      }
+      const totalTokens = this.calculateCallbackTokens(tokenCB, this['request.model'], content)
+      this.setTokensOnEmbeddingMessage(totalTokens)
       return
     }
 
-    const totalTokens = this.getTotalTokens(response)
-    if (totalTokens) {
-      this['response.usage.total_tokens'] = Number(totalTokens)
-    }
-  }
-
-  getTotalTokens(response) {
-    return response?.usage?.total_tokens || response?.usage?.totalTokens
+    const { totalTokens } = this.getUsageTokens(response)
+    this.setTokensOnEmbeddingMessage(totalTokens)
   }
 }

package/lib/llm-events/openai/event.js

@@ -47,4 +47,11 @@ module.exports = class LlmEvent extends BaseEvent {
     this['response.headers.ratelimitRemainingRequests'] =
       response?.headers?.['x-ratelimit-remaining-requests']
   }
+
+  getUsageTokens(response) {
+    const promptTokens = Number(response?.usage?.prompt_tokens || response?.usage?.input_tokens)
+    const completionTokens = Number(response?.usage?.completion_tokens || response?.usage?.output_tokens)
+    const totalTokens = Number(response?.usage?.total_tokens || response?.usage?.totalTokens)
+    return { promptTokens, completionTokens, totalTokens }
+  }
 }

package/lib/shim/datastore-shim.js

@@ -5,7 +5,6 @@
 
 'use strict'
 
-const dbutil = require('../db/utils')
 const hasOwnProperty = require('../util/properties').hasOwn
 const logger = require('../logger').child({ component: 'DatastoreShim' })
 const metrics = require('../metrics/names')
@@ -498,15 +497,14 @@ function captureInstanceAttributes(host, port, database) {
 }
 
 /**
- * Parses the database name from a `USE` SQL query.
+ * Deprecated with mysql subscriber migration.
+ * Logic for mysql database name extraction is now
+ * in `lib/subscribers/mysql/helper.js`.
  *
  * @this DatastoreShim
- * @param {string} query - The SQL query to parse the database name from.
- * @returns {?string} The name of the database if it could be parsed, otherwise
- *  `null`.
  */
-function getDatabaseNameFromUseQuery(query) {
-  return dbutil.extractDatabaseChangeFromUse(query)
+function getDatabaseNameFromUseQuery() {
+  logger.warn('`shim.getDatabaseNameFromUseQuery` has been deprecated, if you were relying on this, you must extract table from `use` yourself.')
 }
 
 // -------------------------------------------------------------------------- //

package/lib/subscriber-configs.js

@@ -15,6 +15,8 @@ const subscribers = {
   ...require('./subscribers/fastify/config'),
   ...require('./subscribers/ioredis/config'),
   ...require('./subscribers/mcp-sdk/config'),
+  ...require('./subscribers/mysql/config'),
+  ...require('./subscribers/mysql2/config'),
   ...require('./subscribers/openai/config'),
   ...require('./subscribers/pino/config'),
   ...require('./subscribers/pg/config'),

package/lib/subscribers/mysql/config.js

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+const connectionQuery = {
+  path: './mysql/connection-query.js',
+  instrumentations: [
+    {
+      channelName: 'nr_connectionQuery',
+      module: { name: 'mysql', versionRange: '>=2.16.0', filePath: 'lib/Connection.js' },
+      functionQuery: {
+        moduleName: 'Connection',
+        expressionName: 'query',
+        kind: 'Sync'
+      }
+    }
+  ]
+}
+
+const poolQuery = {
+  path: './mysql/pool-query.js',
+  instrumentations: [
+    {
+      channelName: 'nr_poolQuery',
+      module: { name: 'mysql', versionRange: '>=2.16.0', filePath: 'lib/Pool.js' },
+      functionQuery: {
+        moduleName: 'Pool',
+        expressionName: 'query',
+        kind: 'Sync'
+      }
+    }
+  ]
+}
+
+const poolNamespaceQuery = {
+  path: './mysql/pool-namespace-query.js',
+  instrumentations: [
+    {
+      channelName: 'nr_poolNamespaceQuery',
+      module: { name: 'mysql', versionRange: '>=2.16.0', filePath: 'lib/PoolNamespace.js' },
+      functionQuery: {
+        moduleName: 'PoolNamespace',
+        expressionName: 'query',
+        kind: 'Sync'
+      }
+    }
+  ]
+}
+
+const poolGetConnection = {
+  path: './mysql/pool-get-connection.js',
+  instrumentations: [
+    {
+      channelName: 'nr_poolGetConnection',
+      module: { name: 'mysql', versionRange: '>=2.16.0', filePath: 'lib/Pool.js' },
+      functionQuery: {
+        moduleName: 'Pool',
+        expressionName: 'getConnection',
+        kind: 'Sync'
+      }
+    }
+  ]
+}
+
+module.exports = {
+  mysql: [
+    connectionQuery,
+    poolQuery,
+    poolNamespaceQuery,
+    poolGetConnection,
+  ]
+}