newrelic 13.6.2 → 13.6.4

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. package/NEWS.md +33 -7
  2. package/README.md +6 -11
  3. package/lib/config/default.js +46 -4
  4. package/lib/config/index.js +110 -1
  5. package/lib/instrumentation/@google/genai.js +15 -24
  6. package/lib/instrumentation/aws-sdk/v3/bedrock.js +5 -0
  7. package/lib/instrumentations.js +0 -1
  8. package/lib/llm-events/aws-bedrock/bedrock-response.js +25 -1
  9. package/lib/llm-events/aws-bedrock/chat-completion-message.js +2 -10
  10. package/lib/llm-events/aws-bedrock/chat-completion-summary.js +11 -15
  11. package/lib/llm-events/aws-bedrock/converse-stream-handler.js +11 -4
  12. package/lib/llm-events/aws-bedrock/embedding.js +3 -20
  13. package/lib/llm-events/aws-bedrock/event.js +3 -20
  14. package/lib/llm-events/aws-bedrock/stream-handler.js +21 -28
  15. package/lib/llm-events/aws-bedrock/utils.js +1 -59
  16. package/lib/llm-events/event.js +108 -0
  17. package/lib/llm-events/google-genai/chat-completion-message.js +3 -10
  18. package/lib/llm-events/google-genai/chat-completion-summary.js +4 -5
  19. package/lib/llm-events/google-genai/embedding.js +0 -34
  20. package/lib/llm-events/google-genai/event.js +7 -0
  21. package/lib/llm-events/openai/chat-completion-message.js +3 -9
  22. package/lib/llm-events/openai/chat-completion-summary.js +4 -5
  23. package/lib/llm-events/openai/embedding.js +4 -15
  24. package/lib/llm-events/openai/event.js +7 -0
  25. package/lib/shim/datastore-shim.js +5 -7
  26. package/lib/subscriber-configs.js +2 -0
  27. package/lib/subscribers/base.js +1 -1
  28. package/lib/subscribers/message-consumer.js +9 -2
  29. package/lib/subscribers/mysql/config.js +73 -0
  30. package/lib/subscribers/mysql/connection-query.js +56 -0
  31. package/lib/subscribers/mysql/helper.js +102 -0
  32. package/lib/subscribers/mysql/pool-get-connection.js +20 -0
  33. package/lib/subscribers/mysql/pool-namespace-query.js +19 -0
  34. package/lib/subscribers/mysql/pool-query.js +41 -0
  35. package/lib/subscribers/mysql2/config.js +127 -0
  36. package/lib/subscribers/mysql2/connection-execute.js +18 -0
  37. package/lib/subscribers/mysql2/connection-query.js +17 -0
  38. package/lib/subscribers/mysql2/pool-get-connection.js +20 -0
  39. package/lib/subscribers/mysql2/pool-namespace-query.js +19 -0
  40. package/lib/subscribers/mysql2/pool-query.js +19 -0
  41. package/lib/subscribers/openai/utils.js +11 -2
  42. package/lib/subscribers/pg/query.js +1 -1
  43. package/package.json +1 -1
  44. package/lib/db/utils.js +0 -19
  45. package/lib/instrumentation/mysql/mysql.js +0 -333
  46. package/lib/instrumentation/mysql/nr-hooks.js +0 -26
  47. package/lib/llm-events/google-genai/utils.js +0 -32
  48. package/lib/llm-events/openai/utils.js +0 -34
  49. package/lib/llm-events/utils.js +0 -110
@@ -6,7 +6,7 @@
6
6
  'use strict'
7
7
 
8
8
  const { randomUUID } = require('crypto')
9
- const { DESTINATIONS } = require('../../config/attribute-filter')
9
+ const BaseEvent = require('../event')
10
10
 
11
11
  /**
12
12
  * @typedef {object} LlmEventParams
@@ -29,7 +29,7 @@ const defaultParams = {
29
29
  /**
30
30
  * Baseline object representing a LLM event.
31
31
  */
32
- class LlmEvent {
32
+ class LlmEvent extends BaseEvent {
33
33
  /**
34
34
  * All parameters that were passed in to the constructor after they have
35
35
  * been merged with the constructor's defaults.
@@ -43,6 +43,7 @@ class LlmEvent {
43
43
  * @param {LlmEventParams} params Construction parameters.
44
44
  */
45
45
  constructor(params = defaultParams) {
46
+ super()
46
47
  params = Object.assign({}, defaultParams, params)
47
48
  this.constructionParams = params
48
49
 
@@ -64,24 +65,6 @@ class LlmEvent {
64
65
  this['request.max_tokens'] = null
65
66
  }
66
67
 
67
- /**
68
- * Pull user set `llm.*` attributes from the current transaction and
69
- * add them to the event.
70
- *
71
- * @param {object} agent The New Relic agent that provides access to the
72
- * transaction.
73
- */
74
- set metadata(agent) { // eslint-disable-line accessor-pairs
75
- const tx = agent.tracer.getTransaction()
76
- const attrs = tx?.trace?.custom.get(DESTINATIONS.TRANS_SCOPE) || {}
77
- for (const [k, v] of Object.entries(attrs)) {
78
- if (k.startsWith('llm.') === false) {
79
- continue
80
- }
81
- this[k] = v
82
- }
83
- }
84
-
85
68
  /**
86
69
  * Removes the complex objects from the event
87
70
  * This will be called right before the event is enqueued to the custom event aggregator
@@ -187,6 +187,13 @@ class StreamHandler {
187
187
  this.response.response.headers = {
188
188
  'x-amzn-requestid': this.passThroughParams.response.response.headers['x-amzn-requestid']
189
189
  }
190
+
191
+ if (parsedEvent['amazon-bedrock-invocationMetrics']) {
192
+ const invocationMetrics = parsedEvent['amazon-bedrock-invocationMetrics']
193
+ this.response.response.headers['x-amzn-bedrock-input-token-count'] = invocationMetrics?.inputTokenCount
194
+ this.response.response.headers['x-amzn-bedrock-output-token-count'] = invocationMetrics?.outputTokenCount
195
+ }
196
+
190
197
  delete parsedEvent['amazon-bedrock-invocationMetrics']
191
198
  }
192
199
  }
@@ -213,8 +220,7 @@ async function * handleClaude() {
213
220
  async function * handleClaude3() {
214
221
  let currentBody = {}
215
222
  let stopReason
216
- const completions = []
217
- const usage = {}
223
+ let response = ''
218
224
 
219
225
  try {
220
226
  for await (const event of this.stream) {
@@ -223,25 +229,14 @@ async function * handleClaude3() {
223
229
  this.updateHeaders(parsed)
224
230
  currentBody = parsed
225
231
  if (parsed.type === 'content_block_delta') {
226
- completions.push(parsed.delta.text)
232
+ response += parsed.delta.text
227
233
  } else if (parsed.type === 'message_delta') {
228
234
  stopReason = parsed.delta.stop_reason
229
235
  }
230
-
231
- if (parsed?.message?.usage?.input_tokens) {
232
- usage.input_tokens = parsed.message.usage.input_tokens
233
- }
234
-
235
- if (parsed?.message?.usage?.output_tokens) {
236
- usage.output_tokens = parsed.message.usage.output_tokens
237
- } else if (parsed?.usage?.output_tokens) {
238
- usage.output_tokens = parsed.usage.output_tokens
239
- }
240
236
  }
241
237
  } finally {
242
- currentBody.completions = completions
238
+ currentBody.completions = response
243
239
  currentBody.stop_reason = stopReason
244
- currentBody.usage = usage
245
240
  this.response.output.body = currentBody
246
241
  this.finish()
247
242
  }
@@ -286,7 +281,6 @@ async function * handleCohereEmbed() {
286
281
  async function * handleLlama() {
287
282
  let currentBody = {}
288
283
  let generation = ''
289
- const usage = {}
290
284
 
291
285
  try {
292
286
  for await (const event of this.stream) {
@@ -295,35 +289,34 @@ async function * handleLlama() {
295
289
  this.updateHeaders(parsed)
296
290
  currentBody = parsed
297
291
  generation += parsed.generation
298
-
299
- if (parsed?.generation_token_count) {
300
- usage.output_tokens = parsed.generation_token_count
301
- }
302
-
303
- if (parsed?.prompt_token_count) {
304
- usage.input_tokens = parsed.prompt_token_count
305
- }
306
292
  }
307
293
  } finally {
308
294
  currentBody.generation = generation
309
- currentBody.usage = usage
310
295
  this.response.output.body = currentBody
311
296
  this.finish()
312
297
  }
313
298
  }
314
299
 
300
+ /**
301
+ * Yields every chunk and builds up the response in a string
302
+ * Re-assigns the response to `outputText` and makes body.results an array
303
+ * to match the structure of non-streamed Titan response.
304
+ */
315
305
  async function * handleTitan() {
316
- const body = this.response.output.body
317
- body.results = []
306
+ let currentBody = {}
307
+ let response = ''
318
308
 
319
309
  try {
320
310
  for await (const event of this.stream) {
321
311
  yield event // Pass it up to the real consumer of the stream.
322
312
  const parsed = this.parseEvent(event)
323
313
  this.updateHeaders(parsed)
324
- body.results.push(parsed)
314
+ response += parsed.outputText
315
+ currentBody = parsed
325
316
  }
326
317
  } finally {
318
+ currentBody.outputText = response
319
+ this.response.output.body.results = [currentBody]
327
320
  this.finish()
328
321
  }
329
322
  }
@@ -5,8 +5,6 @@
5
5
 
6
6
  'use strict'
7
7
 
8
- const { setTokensInResponse } = require('../utils')
9
-
10
8
  /**
11
9
  *
12
10
  * @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) {
69
67
  return stringifiedChunks.join('\n\n')
70
68
  }
71
69
 
72
- /**
73
- * Set the prompt, completion and total token counts on the Llm Event if input and output metadata
74
- * exists in the response
75
- * @param {object} response - The response object.
76
- * @param {object} context - The Llm Event
77
- */
78
- function setUsageTokens(response, context) {
79
- // input and output token counts must available in order to add all usage attributes to response
80
- // if total tokens is not available, we can manually add it up (from input and output token count)
81
- if (tokenUsageAttributesExist(response)) {
82
- const promptTokens =
83
- Number(response?.usage?.input_tokens ||
84
- response?.usage?.inputTokens ||
85
- response?.parsedBody?.results?.[0]?.inputTextTokenCount ||
86
- response?.parsedBody?.usage?.input_tokens)
87
- const completionTokens =
88
- Number(response?.usage?.output_tokens ||
89
- response?.usage?.outputTokens ||
90
- response?.parsedBody?.results?.[1]?.totalOutputTextTokenCount ||
91
- response?.parsedBody?.usage.output_tokens)
92
- const totalTokens = Number(response?.usage?.total_tokens || response?.usage?.totalTokens)
93
-
94
- setTokensInResponse(context, { promptTokens, completionTokens, totalTokens })
95
- return
96
- }
97
-
98
- if (tokenUsageHeadersExist(response)) {
99
- const promptTokens = Number(response?.headers['x-amzn-bedrock-input-token-count'])
100
- const completionTokens = Number(response?.headers['x-amzn-bedrock-output-token-count'])
101
- const totalTokens = Number(response?.headers['x-amzn-bedrock-total-token-count'])
102
-
103
- setTokensInResponse(context, { promptTokens, completionTokens, totalTokens })
104
- }
105
- }
106
-
107
- function tokenUsageAttributesExist(response) {
108
- const tokensA = response?.usage?.input_tokens && response?.usage?.output_tokens
109
- const tokensB = response?.usage?.inputTokens && response?.usage?.outputTokens
110
-
111
- // titan stream puts tokens in here
112
- const tokensC = response?.parsedBody?.results?.[0]?.inputTextTokenCount && response?.parsedBody?.results?.[1]?.totalOutputTextTokenCount
113
-
114
- // for llama3 and claude 3 streaming
115
- const tokensD = response?.parsedBody?.usage?.input_tokens && response?.parsedBody?.usage?.output_tokens
116
-
117
- return tokensA !== undefined || tokensB !== undefined || tokensC !== undefined || tokensD !== undefined
118
- }
119
-
120
- function tokenUsageHeadersExist(response) {
121
- const tokens = response?.headers['x-amzn-bedrock-input-token-count'] && response?.headers['x-amzn-bedrock-output-token-count']
122
- return tokens !== undefined
123
- }
124
-
125
70
  module.exports = {
126
71
  stringifyClaudeChunkedMessage,
127
- stringifyConverseChunkedMessage,
128
- tokenUsageHeadersExist,
129
- tokenUsageAttributesExist,
130
- setUsageTokens,
72
+ stringifyConverseChunkedMessage
131
73
  }
@@ -18,6 +18,114 @@ class BaseLlmEvent {
18
18
  }
19
19
  }
20
20
  }
21
+
22
+ /**
23
+ * Determines if the provided token count is valid.
24
+ * A valid token count is greater than 0 and not null.
25
+ * @param {number} tokenCount The token count obtained from the token callback
26
+ * @returns {boolean} Whether the token count is valid
27
+ */
28
+ validTokenCount(tokenCount) {
29
+ return tokenCount !== null && tokenCount > 0
30
+ }
31
+
32
+ /**
33
+ * Calculates the total token count from the prompt tokens and completion tokens
34
+ * set in the event.
35
+ * @returns {number} The total token count
36
+ */
37
+ getTotalTokenCount() {
38
+ return Number(this['response.usage.prompt_tokens']) + Number(this['response.usage.completion_tokens'])
39
+ }
40
+
41
+ setTokensOnEmbeddingMessage(totalTokens) {
42
+ if (this.validTokenCount(totalTokens)) {
43
+ this['response.usage.total_tokens'] = totalTokens
44
+ }
45
+ }
46
+
47
+ /**
48
+ * Sets the provided tokens counts on the LLM event.
49
+ * Checks if promptToken and completionToken are greater than zero before setting.
50
+ * This is because the spec states that token counts should only be set if both
51
+ * are present.
52
+ * @param {object} params to the function
53
+ * @param {object} params.promptTokens value of prompt token count
54
+ * @param {object} params.completionTokens value of completion(s) token count
55
+ * @param {object} params.totalTokens value of prompt + completion(s) token count
56
+ */
57
+ setTokensInResponse({ promptTokens, completionTokens, totalTokens }) {
58
+ if (this.validTokenCount(promptTokens) && this.validTokenCount(completionTokens)) {
59
+ this['response.usage.prompt_tokens'] = promptTokens
60
+ this['response.usage.completion_tokens'] = completionTokens
61
+ this['response.usage.total_tokens'] = totalTokens || this.getTotalTokenCount()
62
+ }
63
+ }
64
+
65
+ /**
66
+ * Sets `token_count` to 0 on the LlmChatCompletionMessage if both prompt and completion tokens are greater than zero.
67
+ * This is because the spec states that if token counts are set, then we should set token_count to 0 to indicate
68
+ * that the token calculation does not have to occur in the ingest pipeline.
69
+ * @param {object} params to the function
70
+ * @param {object} params.promptTokens value of prompt token count
71
+ * @param {object} params.completionTokens value of completion(s) token count
72
+ */
73
+ setTokenInCompletionMessage({ promptTokens, completionTokens }) {
74
+ if (this.validTokenCount(promptTokens) && this.validTokenCount(completionTokens)) {
75
+ this.token_count = 0
76
+ }
77
+ }
78
+
79
+ /**
80
+ * Calculates prompt and completion token counts using the provided callback and models.
81
+ * If both counts are valid, sets this.token_count to 0.
82
+ *
83
+ * @param {object} options - The params object.
84
+ * @param {Function} options.tokenCB - The token counting callback function.
85
+ * @param {string} options.reqModel - The model used for the prompt.
86
+ * @param {string} options.resModel - The model used for the completion.
87
+ * @param {string} options.promptContent - The prompt content to count tokens for.
88
+ * @param {string} options.completionContent - The completion content to count tokens for.
89
+ * @returns {void}
90
+ */
91
+ setTokenFromCallback({ tokenCB, reqModel, resModel, promptContent, completionContent }) {
92
+ const promptToken = this.calculateCallbackTokens(tokenCB, reqModel, promptContent)
93
+ const completionToken = this.calculateCallbackTokens(tokenCB, resModel, completionContent)
94
+
95
+ this.setTokenInCompletionMessage({ promptTokens: promptToken, completionTokens: completionToken })
96
+ }
97
+
98
+ /**
99
+ * Calculates prompt and completion token counts using the provided callback and models.
100
+ * If both counts are valid, sets token prompt, completion and total counts on the event.
101
+ *
102
+ * @param {object} options - The params object.
103
+ * @param {Function} options.tokenCB - The token counting callback function.
104
+ * @param {string} options.reqModel - The model used for the prompt.
105
+ * @param {string} options.resModel - The model used for the completion.
106
+ * @param {string} options.promptContent - The prompt content to count tokens for.
107
+ * @param {string} options.completionContent - The completion content to count tokens for.
108
+ * @returns {void}
109
+ */
110
+ setTokenUsageFromCallback({ tokenCB, reqModel, resModel, promptContent, completionContent }) {
111
+ const promptTokens = this.calculateCallbackTokens(tokenCB, reqModel, promptContent)
112
+ const completionTokens = this.calculateCallbackTokens(tokenCB, resModel, completionContent)
113
+ this.setTokensInResponse({ promptTokens, completionTokens, totalTokens: promptTokens + completionTokens })
114
+ }
115
+
116
+ /**
117
+ * Calculate the token counts using the provided callback.
118
+ * @param {Function} tokenCB - The token count callback function.
119
+ * @param {string} model - The model.
120
+ * @param {string} content - The content to calculate tokens for, such as prompt or completion response.
121
+ * @returns {number|undefined} - The calculated token count or undefined if callback is not a function.
122
+ */
123
+ calculateCallbackTokens(tokenCB, model, content) {
124
+ if (typeof tokenCB === 'function') {
125
+ return tokenCB(model, content)
126
+ }
127
+ return undefined
128
+ }
21
129
  }
22
130
 
23
131
  module.exports = BaseLlmEvent
@@ -6,8 +6,6 @@
6
6
  'use strict'
7
7
  const LlmEvent = require('./event')
8
8
  const { makeId } = require('../../util/hashes')
9
- const { tokenUsageAttributesExist } = require('./utils')
10
- const { setTokenFromCallback } = require('../utils')
11
9
 
12
10
  module.exports = class LlmChatCompletionMessage extends LlmEvent {
13
11
  constructor({
@@ -50,9 +48,8 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
50
48
  const completionContent = responseContent?.map((content) => content.text).join(' ')
51
49
 
52
50
  if (promptContent && completionContent) {
53
- setTokenFromCallback(
51
+ this.setTokenFromCallback(
54
52
  {
55
- context: this,
56
53
  tokenCB,
57
54
  reqModel: request.model,
58
55
  resModel: this['response.model'],
@@ -64,11 +61,7 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
64
61
  return
65
62
  }
66
63
 
67
- // If no token count callback is available, we need to check the response object
68
- // for usage information and set token_count to 0 if all usage attributes are present.
69
- // Response headers won't have token usage information
70
- if (tokenUsageAttributesExist(response)) {
71
- this.token_count = 0
72
- }
64
+ const tokens = this.getUsageTokens(response)
65
+ this.setTokenInCompletionMessage(tokens)
73
66
  }
74
67
  }
@@ -5,8 +5,6 @@
5
5
 
6
6
  'use strict'
7
7
  const LlmEvent = require('./event')
8
- const { setUsageTokens } = require('./utils')
9
- const { setTokenUsageFromCallback } = require('../utils')
10
8
 
11
9
  module.exports = class LlmChatCompletionSummary extends LlmEvent {
12
10
  constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -39,9 +37,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
39
37
  const completionContent = responseContent?.map((content) => content.text).join(' ')
40
38
 
41
39
  if (promptContent && completionContent) {
42
- setTokenUsageFromCallback(
40
+ this.setTokenUsageFromCallback(
43
41
  {
44
- context: this,
45
42
  tokenCB,
46
43
  reqModel: request.model,
47
44
  resModel: this['response.model'],
@@ -52,6 +49,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
52
49
  }
53
50
  return
54
51
  }
55
- setUsageTokens(response, this)
52
+
53
+ const tokens = this.getUsageTokens(response)
54
+ this.setTokensInResponse(tokens)
56
55
  }
57
56
  }
@@ -6,7 +6,6 @@
6
6
  'use strict'
7
7
 
8
8
  const LlmEvent = require('./event')
9
- const { validCallbackTokenCount, calculateCallbackTokens } = require('../utils')
10
9
 
11
10
  class LlmEmbedding extends LlmEvent {
12
11
  constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -16,39 +15,6 @@ class LlmEmbedding extends LlmEvent {
16
15
  if (agent.config.ai_monitoring.record_content.enabled === true) {
17
16
  this.input = request.contents?.toString()
18
17
  }
19
-
20
- this.setTotalTokens(agent, request, response)
21
- }
22
-
23
- setTotalTokens(agent, request, response) {
24
- const tokenCB = agent.llm?.tokenCountCallback
25
-
26
- // For embedding events, only total token count is relevant.
27
- // Prefer callback for total tokens; if unavailable, fall back to response data.
28
- if (tokenCB) {
29
- const content = request.contents?.toString()
30
-
31
- if (content === undefined) {
32
- return
33
- }
34
-
35
- const totalTokenCount = calculateCallbackTokens(tokenCB, this['request.model'], content)
36
- const hasValidCallbackCounts = validCallbackTokenCount(totalTokenCount)
37
-
38
- if (hasValidCallbackCounts) {
39
- this['response.usage.total_tokens'] = Number(totalTokenCount)
40
- }
41
- return
42
- }
43
-
44
- const totalTokens = this.getTotalTokens(response)
45
- if (totalTokens) {
46
- this['response.usage.total_tokens'] = Number(totalTokens)
47
- }
48
- }
49
-
50
- getTotalTokens(response) {
51
- return response?.usageMetadata?.totalTokenCount
52
18
  }
53
19
  }
54
20
 
@@ -26,4 +26,11 @@ module.exports = class LlmEvent extends BaseEvent {
26
26
  this.duration = segment?.getDurationInMillis()
27
27
  }
28
28
  }
29
+
30
+ getUsageTokens(response) {
31
+ const promptTokens = Number(response?.usageMetadata?.promptTokenCount)
32
+ const completionTokens = Number(response?.usageMetadata?.candidatesTokenCount)
33
+ const totalTokens = Number(response?.usageMetadata?.totalTokenCount)
34
+ return { promptTokens, completionTokens, totalTokens }
35
+ }
29
36
  }
@@ -5,8 +5,6 @@
5
5
 
6
6
  'use strict'
7
7
  const LlmEvent = require('./event')
8
- const { tokenUsageAttributesExist } = require('./utils')
9
- const { setTokenFromCallback } = require('../utils')
10
8
 
11
9
  module.exports = class LlmChatCompletionMessage extends LlmEvent {
12
10
  constructor({
@@ -58,9 +56,8 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
58
56
  : response?.choices?.map((resContent) => resContent.message.content).join(' ')
59
57
 
60
58
  if (promptContent && completionContent) {
61
- setTokenFromCallback(
59
+ this.setTokenFromCallback(
62
60
  {
63
- context: this,
64
61
  tokenCB,
65
62
  reqModel: request.model,
66
63
  resModel: this['response.model'],
@@ -72,10 +69,7 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
72
69
  return
73
70
  }
74
71
 
75
- // If no token count callback is available, we need to check the response object
76
- // for usage information and set token_count to 0 if all usage attributes are present.
77
- if (tokenUsageAttributesExist(response)) {
78
- this.token_count = 0
79
- }
72
+ const tokens = this.getUsageTokens(response)
73
+ this.setTokenInCompletionMessage(tokens)
80
74
  }
81
75
  }
@@ -5,8 +5,6 @@
5
5
 
6
6
  'use strict'
7
7
  const LlmEvent = require('./event')
8
- const { setUsageTokens } = require('./utils')
9
- const { setTokenUsageFromCallback } = require('../utils')
10
8
 
11
9
  module.exports = class LlmChatCompletionSummary extends LlmEvent {
12
10
  constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -46,9 +44,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
46
44
  : response?.choices?.map((resContent) => resContent.message.content).join(' ')
47
45
 
48
46
  if (promptContent && completionContent) {
49
- setTokenUsageFromCallback(
47
+ this.setTokenUsageFromCallback(
50
48
  {
51
- context: this,
52
49
  tokenCB,
53
50
  reqModel: request.model,
54
51
  resModel: this['response.model'],
@@ -59,6 +56,8 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
59
56
  }
60
57
  return
61
58
  }
62
- setUsageTokens(response, this)
59
+
60
+ const tokens = this.getUsageTokens(response)
61
+ this.setTokensInResponse(tokens)
63
62
  }
64
63
  }
@@ -5,7 +5,6 @@
5
5
 
6
6
  'use strict'
7
7
  const LlmEvent = require('./event')
8
- const { validCallbackTokenCount, calculateCallbackTokens } = require('../utils')
9
8
 
10
9
  module.exports = class LlmEmbedding extends LlmEvent {
11
10
  constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -31,22 +30,12 @@ module.exports = class LlmEmbedding extends LlmEvent {
31
30
  return
32
31
  }
33
32
 
34
- const totalTokenCount = calculateCallbackTokens(tokenCB, this['request.model'], content)
35
- const hasValidCallbackCounts = validCallbackTokenCount(totalTokenCount)
36
-
37
- if (hasValidCallbackCounts) {
38
- this['response.usage.total_tokens'] = Number(totalTokenCount)
39
- }
33
+ const totalTokens = this.calculateCallbackTokens(tokenCB, this['request.model'], content)
34
+ this.setTokensOnEmbeddingMessage(totalTokens)
40
35
  return
41
36
  }
42
37
 
43
- const totalTokens = this.getTotalTokens(response)
44
- if (totalTokens) {
45
- this['response.usage.total_tokens'] = Number(totalTokens)
46
- }
47
- }
48
-
49
- getTotalTokens(response) {
50
- return response?.usage?.total_tokens || response?.usage?.totalTokens
38
+ const { totalTokens } = this.getUsageTokens(response)
39
+ this.setTokensOnEmbeddingMessage(totalTokens)
51
40
  }
52
41
  }
@@ -47,4 +47,11 @@ module.exports = class LlmEvent extends BaseEvent {
47
47
  this['response.headers.ratelimitRemainingRequests'] =
48
48
  response?.headers?.['x-ratelimit-remaining-requests']
49
49
  }
50
+
51
+ getUsageTokens(response) {
52
+ const promptTokens = Number(response?.usage?.prompt_tokens || response?.usage?.input_tokens)
53
+ const completionTokens = Number(response?.usage?.completion_tokens || response?.usage?.output_tokens)
54
+ const totalTokens = Number(response?.usage?.total_tokens || response?.usage?.totalTokens)
55
+ return { promptTokens, completionTokens, totalTokens }
56
+ }
50
57
  }
@@ -5,7 +5,6 @@
5
5
 
6
6
  'use strict'
7
7
 
8
- const dbutil = require('../db/utils')
9
8
  const hasOwnProperty = require('../util/properties').hasOwn
10
9
  const logger = require('../logger').child({ component: 'DatastoreShim' })
11
10
  const metrics = require('../metrics/names')
@@ -498,15 +497,14 @@ function captureInstanceAttributes(host, port, database) {
498
497
  }
499
498
 
500
499
  /**
501
- * Parses the database name from a `USE` SQL query.
500
+ * Deprecated with mysql subscriber migration.
501
+ * Logic for mysql database name extraction is now
502
+ * in `lib/subscribers/mysql/helper.js`.
502
503
  *
503
504
  * @this DatastoreShim
504
- * @param {string} query - The SQL query to parse the database name from.
505
- * @returns {?string} The name of the database if it could be parsed, otherwise
506
- * `null`.
507
505
  */
508
- function getDatabaseNameFromUseQuery(query) {
509
- return dbutil.extractDatabaseChangeFromUse(query)
506
+ function getDatabaseNameFromUseQuery() {
507
+ logger.warn('`shim.getDatabaseNameFromUseQuery` has been deprecated, if you were relying on this, you must extract table from `use` yourself.')
510
508
  }
511
509
 
512
510
  // -------------------------------------------------------------------------- //
@@ -15,6 +15,8 @@ const subscribers = {
15
15
  ...require('./subscribers/fastify/config'),
16
16
  ...require('./subscribers/ioredis/config'),
17
17
  ...require('./subscribers/mcp-sdk/config'),
18
+ ...require('./subscribers/mysql/config'),
19
+ ...require('./subscribers/mysql2/config'),
18
20
  ...require('./subscribers/openai/config'),
19
21
  ...require('./subscribers/pino/config'),
20
22
  ...require('./subscribers/pg/config'),
@@ -363,7 +363,7 @@ class Subscriber {
363
363
  if (distributedTracingEnabled) {
364
364
  // we have to pass in traceId, segment id, and hard code traceFlags to 1
365
365
  // because we're inserting headers right before the original function is bound.
366
- const traceFlags = tx.sampled === true ? 1 : 0
366
+ const traceFlags = tx.isSampled() === true ? 1 : 0
367
367
  tx.insertDistributedTraceHeaders(headers, null, { traceId: tx.traceId, spanId: ctx?.segment?.id, traceFlags })
368
368
  } else {
369
369
  cat.addCatHeaders(this.config, tx, headers, useMqNames)
@@ -65,6 +65,14 @@ class MessageConsumerSubscriber extends Subscriber {
65
65
  }
66
66
  }
67
67
 
68
+ /**
69
+ * Ends the transaction created for the consumption callback.
70
+ */
71
+ asyncEnd() {
72
+ const ctx = this.agent.tracer.getContext()
73
+ ctx?.transaction?.end()
74
+ }
75
+
68
76
  enable() {
69
77
  super.enable()
70
78
  this.channel.asyncStart.bindStore(this.store, (data) => {
@@ -83,9 +91,8 @@ class MessageConsumerSubscriber extends Subscriber {
83
91
  /**
84
92
  * Used to create a transaction for every consumption callback.
85
93
  *
86
- * @param {object} data event passed to asyncStart handler
87
94
  */
88
- asyncStart(data) {
95
+ asyncStart() {
89
96
  const ctx = this.agent.tracer.getContext()
90
97
  const tx = ctx?.transaction
91
98
  tx.setPartialName(this.name)