aws bedrock embedding

Author: amychisholm03

lib/instrumentation/aws-sdk/v3/bedrock.js

@@ -7,12 +7,14 @@
 const {
   LlmChatCompletionMessage,
   LlmChatCompletionSummary,
-  LlmEmbedding,
   LlmErrorMessage,
   BedrockCommand,
   BedrockResponse,
   StreamHandler
 } = require('../../../llm-events/aws-bedrock')
+const {
+  LlmEmbedding
+} = require('../../../llm-events-new/aws-bedrock')
 
 const { DESTINATIONS } = require('../../../config/attribute-filter')
 const { AI } = require('../../../metrics/names')
@@ -73,7 +75,7 @@ function isStreamingEnabled({ commandName, config }) {
  * @param {object} params.msg LLM event
  */
 function recordEvent({ agent, type, msg }) {
-  msg.serialize()
+  if (msg.serialize) msg.serialize() // TODO: remove
   const llmContext = extractLlmContext(agent)
   const timestamp = msg?.timestamp ?? Date.now()
 
@@ -216,10 +218,11 @@ function recordEmbeddingMessage({
     agent,
     segment,
     transaction,
-    bedrockCommand,
-    input: prompt.content,
-    bedrockResponse,
-    isError: err !== null
+    requestModel: bedrockCommand?.modelId,
+    requestInput: prompt.content,
+    requestId: bedrockResponse?.requestId,
+    totalTokenCount: bedrockResponse?.totalTokenCount,
+    error: err !== null
   }))
 
   for (const embedding of embeddings) {

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2026 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+const LlmEmbedding = require('../embedding')
+
+/**
+ * Encapsulates a AWS Bedrock LlmEmbedding event.
+ */
+module.exports = class AwsBedrockLlmEmbedding extends LlmEmbedding {
+  /**
+   *
+   * @param {object} params constructor parameters
+   * @param {Agent} params.agent New Relic agent instance
+   * @param {object} params.segment Current segment
+   * @param {object} params.transaction Current and active transaction
+   * @param {string} params.requestId ID associated with the request - typically available in response headers
+   * @param {string} params.requestInput Input to the embedding creation call
+   * @param {string} params.requestModel Model name specified in the request (e.g. 'gpt-4')
+   * @param {number} params.totalTokenCount Retrieved from the Bedrock response object, fallback for token calculation
+   * @param {boolean} [params.error] Set to `true` if an error occurred during creation call - omitted if no error occurred
+   */
+  constructor({ agent, segment, transaction, requestInput, requestModel, requestId, totalTokenCount = 0, error }) {
+    super({ agent,
+      segment,
+      transaction,
+      vendor: 'bedrock',
+      requestInput,
+      requestModel,
+      requestId,
+      error })
+
+    this.setTotalTokens(agent, requestInput, totalTokenCount)
+  }
+
+  setTotalTokens(agent, input, totalTokenCount) {
+    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 = input?.toString()
+
+      if (content === undefined) {
+        return
+      }
+
+      const totalTokens = this.calculateCallbackTokens(tokenCB, this['request.model'], content)
+      this.setTokensOnEmbeddingMessage(totalTokens)
+      return
+    }
+
+    this.setTokensOnEmbeddingMessage(totalTokenCount)
+  }
+}

lib/llm-events-new/aws-bedrock/index.js

@@ -0,0 +1,12 @@
+/*
+ * Copyright 2026 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const LlmEmbedding = require('./embedding')
+
+module.exports = {
+  LlmEmbedding
+}

lib/llm-events-new/base.js

@@ -86,6 +86,11 @@ class LlmEvent {
     return Number(this['response.usage.prompt_tokens']) + Number(this['response.usage.completion_tokens'])
   }
 
+  /**
+   * If `totalTokens` is valid, assigns it to
+   * `this['response.usage.total_tokens']`.
+   * @param {number} totalTokens total tokens on embedding message
+   */
   setTokensOnEmbeddingMessage(totalTokens) {
     if (this.validTokenCount(totalTokens)) {
       this['response.usage.total_tokens'] = totalTokens

lib/llm-events-new/embedding.js

@@ -23,11 +23,11 @@ class LlmEmbedding extends LlmEvent {
    * @param {Agent} params.agent New Relic agent instance
    * @param {object} params.segment Current segment
    * @param {object} params.transaction Current and active transaction
-   * @param {string} params.requestId ID associated with the request - typically available in response headers
+   * @param {string} [params.requestId] ID associated with the request - typically available in response headers
    * @param {string} params.requestInput Input to the embedding creation call
-   * @param {string} params.requestModel Model name specified in the request (e.g. 'gpt-4')
-   * @param {string} params.responseModel Model name returned in the response (can differ from `request.model`)
-   * @param {string} params.responseOrg Organization ID returned in the response or response headers
+   * @param {string} [params.requestModel] Model name specified in the request (e.g. 'gpt-4')
+   * @param {string} [params.responseModel] Model name returned in the response (can differ from `request.model`)
+   * @param {string} [params.responseOrg] Organization ID returned in the response or response headers
    * @param {string} params.vendor Lowercased name of vendor (e.g. 'openai')
    * @param {boolean} [params.error] Set to `true` if an error occurred during creation call - omitted if no error occurred
    */

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

@@ -10,7 +10,6 @@ module.exports = {
   BedrockResponse: require('./bedrock-response'),
   LlmChatCompletionMessage: require('./chat-completion-message'),
   LlmChatCompletionSummary: require('./chat-completion-summary'),
-  LlmEmbedding: require('./embedding'),
   LlmEvent: require('./event'),
   LlmErrorMessage: require('../error-message'),
   StreamHandler: require('./stream-handler'),

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

@@ -10,7 +10,7 @@ const assert = require('node:assert')
 const {
   DESTINATIONS: { TRANS_SCOPE }
 } = require('../../../../lib/config/attribute-filter')
-const LlmEmbedding = require('../../../../lib/llm-events/aws-bedrock/embedding')
+const LlmEmbedding = require('#agentlib/llm-events-new/aws-bedrock/embedding.js')
 
 test.beforeEach((ctx) => {
   ctx.nr = {}
@@ -45,19 +45,10 @@ test.beforeEach((ctx) => {
     }
   }
 
-  ctx.nr.bedrockCommand = {
-  }
-
-  ctx.nr.input = 'who are you'
+  ctx.nr.requestInput = 'who are you'
+  ctx.nr.requestId = 'request-1'
+  ctx.nr.totalTokenCount = 70
 
-  ctx.nr.bedrockResponse = {
-    headers: {
-      'x-amzn-requestid': 'request-1'
-    },
-    get inputTokenCount() {
-      return 70
-    }
-  }
   ctx.nr.transaction = {
     traceId: 'id'
   }
@@ -82,17 +73,13 @@ test('should not capture input when `ai_monitoring.record_content.enabled` is fa
   assert.equal(event.input, undefined, 'input should be empty')
 })
 
-test('capture total token usage attribute when inputTokenCount is set', async (t) => {
+test('capture total token usage attribute when totalTokenCount is set', async (t) => {
   const event = new LlmEmbedding(t.nr)
   assert.equal(event['response.usage.total_tokens'], 70)
 })
 
-test('does not capture total token usage when inputTokenCount is not set', async (t) => {
-  Object.defineProperty(t.nr.bedrockResponse, 'inputTokenCount', {
-    get() {
-      return undefined
-    }
-  })
+test('does not capture total token usage when totalTokenCount is not set', async (t) => {
+  t.nr.totalTokenCount = undefined
   const event = new LlmEmbedding(t.nr)
   assert.equal(event['response.usage.total_tokens'], undefined)
 })
@@ -106,14 +93,3 @@ test('should use token callback to set total token usage attribute', async (t) =
 
   assert.equal(event['response.usage.total_tokens'], 65)
 })
-
-test('should not call token callback if there is no content', async (t) => {
-  function cb(model, content) {
-    return 65
-  }
-  t.nr.agent.llm.tokenCountCallback = cb
-  t.nr.input = undefined
-  const event = new LlmEmbedding(t.nr)
-
-  assert.equal(event['response.usage.total_tokens'], undefined)
-})

test/unit/llm-events/aws-bedrock/embedding.test.js

@@ -459,7 +459,7 @@ test('cohere embedding streaming works', async (t) => {
     const events = agent.customEventAggregator.events.toArray()
     assert.equal(events.length, 1)
     const embedding = events.shift()[1]
-    assert.equal(embedding.error, false)
+    assert.equal(embedding.error, undefined)
     assert.equal(embedding.input, prompt)
 
     tx.end()