address feedback

lib/llm/bedrock-response.js

@@ -122,6 +122,25 @@ class BedrockResponse {
     return this.#id
   }
 
+  /**
+   * The number of tokens present in the prompt as determined by the remote
+   * API.
+   *
+   * @returns {number|undefined}
+   */
+  get inputTokenCount() {
+    return this.#tokenCount('x-amzn-bedrock-input-token-count')
+  }
+
+  /**
+   * The number of tokens in the LLM response as determined by the remote API.
+   *
+   * @returns {number|undefined}
+   */
+  get outputTokenCount() {
+    return this.#tokenCount('x-amzn-bedrock-output-token-count')
+  }
+
   /**
    * UUID assigned to the initial request as returned by the API.
    *
@@ -139,6 +158,14 @@ class BedrockResponse {
   get statusCode() {
     return this.#innerResponse.statusCode
   }
+
+  #tokenCount(headerName) {
+    const headerVal = this.headers?.[headerName]
+    if (headerVal != null) {
+      return parseInt(headerVal, 10)
+    }
+    return undefined
+  }
 }
 
 module.exports = BedrockResponse

tests/versioned/common.js

@@ -118,8 +118,15 @@ function assertChatCompletionMessages({ tx, chatMsgs, expectedId, modelId, promp
   })
 }
 
-function assertChatCompletionSummary({ tx, modelId, chatSummary, error = false, numMsgs = 2 }) {
-  const expectedChatSummary = {
+function assertChatCompletionSummary({
+  tx,
+  modelId,
+  chatSummary,
+  tokenUsage,
+  error = false,
+  numMsgs = 2
+}) {
+  let expectedChatSummary = {
     'id': /[\w]{8}-[\w]{4}-[\w]{4}-[\w]{4}-[\w]{12}/,
     'appName': 'New Relic for Node.js tests',
     'request_id': 'eda0760a-c3f0-4fc1-9a1e-75559d642866',
@@ -139,6 +146,15 @@ function assertChatCompletionSummary({ tx, modelId, chatSummary, error = false,
     'error': error
   }
 
+  if (tokenUsage) {
+    expectedChatSummary = {
+      ...expectedChatSummary,
+      'response.usage.total_tokens': 12,
+      'response.usage.prompt_tokens': 8,
+      'response.usage.completion_tokens': 4
+    }
+  }
+
   this.equal(chatSummary[0].type, 'LlmChatCompletionSummary')
   this.match(chatSummary[1], expectedChatSummary, 'should match chat summary message')
 }

tests/versioned/v3/bedrock-embeddings.tap.js

@@ -117,9 +117,9 @@ tap.afterEach(async (t) => {
         'ingest_source': 'Node',
         'request.model': modelId,
         'duration': tx.trace.root.children[0].getDurationInMillis(),
-        'response.usage.total_tokens': undefined,
-        'response.usage.prompt_tokens': undefined,
-        'token_count': undefined,
+        'response.usage.total_tokens': 13,
+        'response.usage.prompt_tokens': 13,
+        'token_count': 13,
         'input': prompt,
         'error': false
       }