feat: Added ability to store feedback ids by request id (newrelic#240)

lib/llm/bedrock-response.js

@@ -136,7 +136,7 @@ class BedrockResponse {
  * @returns {string}
  */
  get requestId() {
- return this.#innerOutput.requestId
+ return this.headers?.['x-amzn-requestid']
  }
 
  /**

lib/llm/event.js

@@ -64,7 +64,7 @@ class LlmEvent {
  this.span_id = segment.id
  this.transaction_id = segment.transaction.id
  this.trace_id = segment.transaction.traceId
- this.request_id = this.bedrockResponse.headers['x-amzn-requestid']
+ this.request_id = this.bedrockResponse.requestId
 
  this['response.model'] = this.bedrockCommand.modelId
  this['request.model'] = this.bedrockCommand.modelId

lib/llm/index.js

@@ -11,5 +11,6 @@ module.exports = {
  LlmChatCompletionMessage: require('./chat-completion-message'),
  LlmChatCompletionSummary: require('./chat-completion-summary'),
  LlmEmbedding: require('./embedding'),
- LlmEvent: require('./event')
+ LlmEvent: require('./event'),
+ LlmTrackedIds: require('./tracked-ids')
 }

lib/llm/tracked-ids.js

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+/**
+ * Represents the identifiers created and collected when creating LLM
+ * chat completions.
+ *
+ * @property {string[]} [message_ids=[]] Internal identifiers for each message
+ * sent in the conversation.
+ * @property {string} [coversation_id=""] User defined identifier of the chat
+ * completion conversation.
+ * @property {string} request_id Identifier of the request from the remote
+ * service for a chat completion. Populated by the `x-request-id` header
+ * in the request.
+ * @public
+ */
+module.exports = class LlmTrackedIds {
+ constructor({ messageIds, conversationId = '', requestId } = {}) {
+ this.message_ids = messageIds ?? []
+ this.conversation_id = conversationId
+ this.request_id = requestId
+ }
+}

lib/v3/bedrock.js

@@ -8,6 +8,7 @@ const {
  LlmChatCompletionMessage,
  LlmChatCompletionSummary,
  LlmEmbedding,
+ LlmTrackedIds,
  BedrockCommand,
  BedrockResponse
 } = require('../llm')
@@ -36,6 +37,28 @@ function recordEvent({ agent, type, msg }) {
  agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
 }
 
+/**
+ * Assigns requestId, conversationId and messageIds for a given
+ * chat completion response on the active transaction.
+ * This is used for generating LlmFeedbackEvent via `api.recordLlmFeedbackEvent`
+ *
+ * @param {object} params input params
+ * @param {Transaction} params.tx active transaction
+ * @param {LlmChatCompletionMessage} params.completionMsg chat completion message
+ * @param {string} params.responseId id of response
+ */
+function assignIdsToTx({ tx, msg, responseId }) {
+ const tracker = tx.llm.responses
+ const trackedIds =
+ tracker.get(responseId) ??
+ new LlmTrackedIds({
+ requestId: msg.request_id,
+ conversationId: msg.conversation_id
+ })
+ trackedIds.message_ids.push(msg.id)
+ tracker.set(responseId, trackedIds)
+}
+
 /**
  * Registers the specification for instrumentation bedrock calls
  *
@@ -76,6 +99,7 @@ function getBedrockSpec({ config, commandName }, _shim, _original, _name, args)
  const bedrockResponse = new BedrockResponse({ bedrockCommand, response })
 
  if (modelType === 'completion') {
+ const tx = segment.transaction
  const summary = new LlmChatCompletionSummary({
  credentials,
  agent,
@@ -92,6 +116,7 @@ function getBedrockSpec({ config, commandName }, _shim, _original, _name, args)
  index: 0,
  completionId: summary.id
  })
+ assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg })
  recordEvent({ agent, type: 'LlmChatCompletionMessage', msg })
 
  bedrockResponse.completions.forEach((content, index) => {
@@ -105,6 +130,7 @@ function getBedrockSpec({ config, commandName }, _shim, _original, _name, args)
  content,
  completionId: summary.id
  })
+ assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg: chatCompletionMessage })
  recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: chatCompletionMessage })
  })
 

tests/unit/llm/bedrock-response.tap.js

@@ -57,13 +57,13 @@ tap.beforeEach((t) => {
  response: {
  statusCode: 200,
  headers: {
+ 'x-amzn-requestid': 'aws-request-1',
  'x-foo': 'foo',
  ['x-amzn-bedrock-input-token-count']: 25,
  ['x-amzn-bedrock-output-token-count']: 25
  }
  },
  output: {
- requestId: 'aws-request-1',
  body: new TextEncoder().encode('{"foo":"foo"}')
  }
  }
@@ -100,7 +100,7 @@ tap.test('non-conforming response is handled gracefully', async (t) => {
  t.equal(res.id, undefined)
  t.equal(res.inputTokenCount, 0)
  t.equal(res.outputTokenCount, 0)
- t.equal(res.requestId, 'aws-request-1')
+ t.equal(res.requestId, undefined)
  t.equal(res.statusCode, 200)
 })
 

tests/unit/llm/event.tap.js

@@ -46,9 +46,7 @@ tap.beforeEach((t) => {
  }
 
  t.context.bedrockResponse = {
- headers: {
- 'x-amzn-requestid': 'request-1'
- }
+ requestId: 'request-1'
  }
 
  t.context.bedrockCommand = {

tests/versioned/common.js

@@ -117,6 +117,7 @@ function assertChatCompletionSummary({ tx, modelId, chatSummary, tokenUsage, err
  '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',
+ 'conversation_id': 'convo-id',
  'trace_id': tx.traceId,
  'span_id': tx.trace.root.children[0].id,
  'transaction_id': tx.id,
@@ -128,6 +129,8 @@ function assertChatCompletionSummary({ tx, modelId, chatSummary, tokenUsage, err
  'api_key_last_four_digits': 'E ID',
  'response.number_of_messages': 2,
  'response.choices.finish_reason': 'endoftext',
+ 'request.temperature': 0.5,
+ 'request.max_tokens': 100,
  'error': error
  }
 

tests/versioned/v3/bedrock-chat-completions.tap.js

@@ -13,23 +13,26 @@ const createAiResponseServer = require('../aws-server-stubs/ai-server')
 const { FAKE_CREDENTIALS } = require('../aws-server-stubs')
 const requests = {
  ai21: (prompt, modelId) => ({
- body: JSON.stringify({ prompt }),
+ body: JSON.stringify({ prompt, temperature: 0.5, maxTokens: 100 }),
  modelId
  }),
  amazon: (prompt, modelId) => ({
- body: JSON.stringify({ inputText: prompt }),
+ body: JSON.stringify({
+ inputText: prompt,
+ textGenerationConfig: { temperature: 0.5, maxTokenCount: 100 }
+ }),
  modelId
  }),
  claude: (prompt, modelId) => ({
- body: JSON.stringify({ prompt }),
+ body: JSON.stringify({ prompt, temperature: 0.5, max_tokens_to_sample: 100 }),
  modelId
  }),
  cohere: (prompt, modelId) => ({
- body: JSON.stringify({ prompt }),
+ body: JSON.stringify({ prompt, temperature: 0.5, max_tokens: 100 }),
  modelId
  }),
  llama2: (prompt, modelId) => ({
- body: JSON.stringify({ prompt }),
+ body: JSON.stringify({ prompt, max_gen_length: 100, temperature: 0.5 }),
  modelId
  })
 }
@@ -116,7 +119,9 @@ tap.afterEach(async (t) => {
  const command = new bedrock.InvokeModelCommand(input)
 
  const { agent } = helper
+ const api = helper.getAgentApi()
  helper.runInTransaction(async (tx) => {
+ api.addCustomAttribute('llm.conversation_id', 'convo-id')
  await client.send(command)
  const events = agent.customEventAggregator.events.toArray()
  t.equal(events.length, 3)
@@ -153,4 +158,54 @@ tap.afterEach(async (t) => {
  t.equal(error.message, expected.body.message)
  }
  })
+
+ tap.test('should store ids and record feedback message accordingly', (t) => {
+ const { bedrock, client, helper } = t.context
+ const conversationId = 'convo-id'
+ const prompt = `text ${resKey} ultimate question`
+ const input = requests[resKey](prompt, modelId)
+ const command = new bedrock.InvokeModelCommand(input)
+
+ const { agent } = helper
+ const api = helper.getAgentApi()
+ helper.runInTransaction(async (tx) => {
+ api.addCustomAttribute('llm.conversation_id', conversationId)
+ const response = await client.send(command)
+ const responseId = response.$metadata.requestId
+ const events = agent.customEventAggregator.events.toArray()
+ const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage')
+ const ids = api.getLlmMessageIds({ responseId })
+ const messageIds = chatMsgs.map((msg) => msg[1].id)
+ t.equal(ids.request_id, responseId)
+ t.equal(ids.conversation_id, conversationId)
+ // message_ids order varies over test run, sort them to assure consistency
+ t.same(ids.message_ids.sort(), messageIds.sort())
+ api.recordLlmFeedbackEvent({
+ conversationId: ids.conversation_id,
+ requestId: ids.request_id,
+ messageId: ids.message_ids[0],
+ category: 'test-event',
+ rating: '5 star',
+ message: 'You are a mathematician.',
+ metadata: { foo: 'foo' }
+ })
+ const recordedEvents = agent.customEventAggregator.getEvents()
+ const [[, feedback]] = recordedEvents.filter(([{ type }]) => type === 'LlmFeedbackMessage')
+
+ t.match(feedback, {
+ id: /[\w\d]{32}/,
+ conversation_id: ids.conversation_id,
+ request_id: ids.request_id,
+ message_id: ids.message_ids[0],
+ category: 'test-event',
+ rating: '5 star',
+ message: 'You are a mathematician.',
+ ingest_source: 'Node',
+ foo: 'foo'
+ })
+
+ tx.end()
+ t.end()
+ })
+ })
 })