fix: Updated AWS Bedrock instrumentation to properly assign token counts (newrelic#3491)

Author: bizob2828

lib/llm-events/aws-bedrock/bedrock-response.js

@@ -73,6 +73,30 @@ class BedrockResponse {
     return this.#parsedBody
   }
 
+  get inputTokenCount() {
+    if (this.#isConverse) {
+      return parseInt(this.#originalResponse?.output?.usage?.inputTokens || 0, 10)
+    }
+
+    return parseInt(this?.headers?.['x-amzn-bedrock-input-token-count'] || 0, 10)
+  }
+
+  get outputTokenCount() {
+    if (this.#isConverse) {
+      return parseInt(this.#originalResponse?.output?.usage?.outputTokens || 0, 10)
+    }
+
+    return parseInt(this?.headers?.['x-amzn-bedrock-output-token-count'] || 0, 10)
+  }
+
+  get totalTokenCount() {
+    if (this.#isConverse) {
+      return parseInt(this.#originalResponse?.output?.usage?.totalTokens || 0, 10)
+    }
+
+    return this.inputTokenCount + this.outputTokenCount
+  }
+
   /**
    * The prompt responses returned by the model.
    *
@@ -164,7 +188,7 @@ class BedrockResponse {
     } else if (cmd.isClaude3() === true) {
       if (body?.type === 'message_stop') {
         // Streamed response
-        this.#completions = body.completions
+        this.#completions.push(body.completions)
       } else {
         this.#completions = [stringifyClaudeChunkedMessage(body?.content)]
       }

lib/llm-events/aws-bedrock/chat-completion-message.js

@@ -6,9 +6,6 @@
 'use strict'
 
 const LlmEvent = require('./event')
-const { setTokenFromCallback } = require('../utils')
-const { tokenUsageAttributesExist, tokenUsageHeadersExist } = require('./utils')
-
 /**
  * @typedef {object} LlmChatCompletionParams
  * @augments LlmEventParams
@@ -71,9 +68,8 @@ class LlmChatCompletionMessage extends LlmEvent {
       const completionContent = this.bedrockResponse?.completions?.join(' ')
 
       if (promptContent && completionContent) {
-        setTokenFromCallback(
+        this.setTokenFromCallback(
           {
-            context: this,
             tokenCB,
             reqModel: this.bedrockCommand.modelId,
             resModel: this.bedrockCommand.modelId,
@@ -85,11 +81,7 @@ class LlmChatCompletionMessage extends LlmEvent {
       return
     }
 
-    // If no token count callback is available, we need to check the response object
-    // or response headers for usage information and set token_count to 0 if all usage attributes are present.
-    if (tokenUsageAttributesExist(this.bedrockResponse) || tokenUsageHeadersExist(this.bedrockResponse)) {
-      this.token_count = 0
-    }
+    this.setTokenInCompletionMessage({ promptTokens: this.bedrockResponse.inputTokenCount, completionTokens: this.bedrockResponse.outputTokenCount })
   }
 }
 

lib/llm-events/aws-bedrock/chat-completion-summary.js

@@ -7,8 +7,6 @@
 
 const LlmEvent = require('./event')
 
-const { setUsageTokens } = require('./utils')
-const { setTokenUsageFromCallback } = require('../utils')
 /**
  * @typedef {object} LlmChatCompletionSummaryParams
  * @augments LlmEventParams
@@ -50,21 +48,19 @@ class LlmChatCompletionSummary extends LlmEvent {
       const promptContent = this.bedrockCommand?.prompt?.map((msg) => msg.content).join(' ')
       const completionContent = this.bedrockResponse?.completions?.join(' ')
 
-      if (promptContent && completionContent) {
-        setTokenUsageFromCallback(
-          {
-            context: this,
-            tokenCB,
-            reqModel: this.bedrockCommand.modelId,
-            resModel: this.bedrockCommand.modelId,
-            promptContent,
-            completionContent
-          }
-        )
-      }
+      this.setTokenUsageFromCallback(
+        {
+          tokenCB,
+          reqModel: this.bedrockCommand.modelId,
+          resModel: this.bedrockCommand.modelId,
+          promptContent,
+          completionContent
+        }
+      )
       return
     }
-    setUsageTokens(this.bedrockResponse, this)
+
+    this.setTokensInResponse({ promptTokens: this.bedrockResponse.inputTokenCount, completionTokens: this.bedrockResponse.outputTokenCount, totalTokens: this.bedrockResponse.totalTokenCount })
   }
 }
 

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

@@ -65,6 +65,11 @@ class ConverseStreamHandler {
         message: {
 
         }
+      },
+      usage: {
+        inputTokens: null,
+        outputTokens: null,
+        totalTokens: null
       }
     }
   }
@@ -119,9 +124,7 @@ async function * handleConverse() {
     for await (const event of this.stream) {
       yield event
       this.updateHeaders(event)
-      if (event.messageStart?.role) {
-        this.role = 'assistant'
-      } else if (event.contentBlockStart?.start) {
+      if (event.contentBlockStart?.start) {
         // Handles a Content block start event. Tool use only.
         const blockStartData = event.contentBlockStart.start
         if (blockStartData.toolUse) {
@@ -133,6 +136,7 @@ async function * handleConverse() {
           activeChunk = null
         }
       } else if (event.contentBlockDelta?.delta) {
+        // There are also deltas for tool use (stringified inputs) but we don't currently record them so we just ignore for now
         if (event.contentBlockDelta.delta.text) {
           // It seems like the first streamed chunk does not always start with a contentBlockStart message
           // If the stream starts with a delta, assume the current chunk is text
@@ -141,9 +145,12 @@ async function * handleConverse() {
           }
           activeChunk.text += event.contentBlockDelta.delta.text
         }
-        // There are also deltas for tool use (stringified inputs) but we don't currently record them so we just ignore for now
+      // used in `handleResponse` to determine why the model stopped
       } else if (event.messageStop) {
         this.stopReason = event.messageStop?.stopReason
+      // sends token usage info at the end of the stream
+      } else if (event.metadata?.usage) {
+        this.response.output.usage = event.metadata.usage
       }
     }
   } finally {

lib/llm-events/aws-bedrock/embedding.js

@@ -6,7 +6,6 @@
 'use strict'
 
 const LlmEvent = require('./event')
-const { validCallbackTokenCount, calculateCallbackTokens } = require('../utils')
 
 /**
  * @typedef {object} LlmEmbeddingParams
@@ -44,28 +43,12 @@ 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()
-    if (totalTokens) {
-      this['response.usage.total_tokens'] = Number(totalTokens)
-    }
-  }
-
-  getTotalTokens() {
-    // We record the input token count as total tokens
-    const totalToken =
-      this.bedrockResponse?.usage?.input_tokens || this.bedrockResponse?.usage?.inputTokens ||
-      this.bedrockResponse?.headers['x-amzn-bedrock-input-token-count']
-
-    return totalToken
+    this.setTokensOnEmbeddingMessage(this.bedrockResponse.inputTokenCount)
   }
 }
 

lib/llm-events/aws-bedrock/event.js

@@ -6,7 +6,7 @@
 'use strict'
 
 const { randomUUID } = require('crypto')
-const { DESTINATIONS } = require('../../config/attribute-filter')
+const BaseEvent = require('../event')
 
 /**
  * @typedef {object} LlmEventParams
@@ -29,7 +29,7 @@ const defaultParams = {
 /**
  * Baseline object representing a LLM event.
  */
-class LlmEvent {
+class LlmEvent extends BaseEvent {
   /**
    * All parameters that were passed in to the constructor after they have
    * been merged with the constructor's defaults.
@@ -43,6 +43,7 @@ class LlmEvent {
    * @param {LlmEventParams} params Construction parameters.
    */
   constructor(params = defaultParams) {
+    super()
     params = Object.assign({}, defaultParams, params)
     this.constructionParams = params
 
@@ -64,24 +65,6 @@ class LlmEvent {
     this['request.max_tokens'] = null
   }
 
-  /**
-   * Pull user set `llm.*` attributes from the current transaction and
-   * add them to the event.
-   *
-   * @param {object} agent The New Relic agent that provides access to the
-   * transaction.
-   */
-  set metadata(agent) { // eslint-disable-line accessor-pairs
-    const tx = agent.tracer.getTransaction()
-    const attrs = tx?.trace?.custom.get(DESTINATIONS.TRANS_SCOPE) || {}
-    for (const [k, v] of Object.entries(attrs)) {
-      if (k.startsWith('llm.') === false) {
-        continue
-      }
-      this[k] = v
-    }
-  }
-
   /**
    * Removes the complex objects from the event
    * This will be called right before the event is enqueued to the custom event aggregator

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()
   }
 }