chore: Refactor LLM user added metadata (newrelic#256)

lib/llm/chat-completion-message.js

@@ -39,10 +39,9 @@ class LlmChatCompletionMessage extends LlmEvent {
     params = Object.assign({}, defaultParams, params)
     super(params)
 
-    const { agent, content, isResponse, index, completionId } = params
+    const { content, isResponse, index, completionId } = params
 
     this.is_response = isResponse
-    this.conversation_id = this.conversationId(agent)
     this.completion_id = completionId
     this.sequence = index
     this.content = content

lib/llm/chat-completion-summary.js

@@ -23,9 +23,8 @@ class LlmChatCompletionSummary extends LlmEvent {
   constructor(params = defaultParams) {
     super(params)
 
-    const { agent, segment, isError } = params
+    const { segment, isError } = params
     this.error = isError
-    this.conversation_id = this.conversationId(agent)
     this.duration = segment.getDurationInMillis()
     this['request.max_tokens'] = this.bedrockCommand.maxTokens
 

lib/llm/event.js

@@ -59,25 +59,29 @@ class LlmEvent {
     this.transaction_id = segment.transaction.id
     this.trace_id = segment.transaction.traceId
     this.request_id = this.bedrockResponse.requestId
+    this.metadata = agent
 
     this['response.model'] = this.bedrockCommand.modelId
     this['request.model'] = this.bedrockCommand.modelId
     this['request.max_tokens'] = null
   }
 
   /**
-   * Retrieve the conversation identifier from the custom attributes
-   * stored in the current transaction.
+   * 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.
-   *
-   * @returns {string}
    */
-  conversationId(agent) {
+  set metadata(agent) {
     const tx = agent.tracer.getTransaction()
-    const attrs = tx?.trace?.custom.get(DESTINATIONS.TRANS_SCOPE)
-    return attrs?.['llm.conversation_id']
+    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
+    }
   }
 
   /**

lib/v3/bedrock.js

@@ -35,33 +35,13 @@ function shouldSkipInstrumentation(config) {
  * @param {Agent} params.agent NR agent instance
  * @param {string} params.type LLM event type
  * @param {object} params.msg LLM event
- * @param {object} params.tx The current transaction.
  */
-function recordEvent({ agent, type, msg, tx }) {
+function recordEvent({ agent, type, msg }) {
   agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount()
-  setMetadata({ tx, msg })
   msg.serialize()
   agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
 }
 
-/**
- * Decorate the outgoing message with metadata that has been set by the
- * customer via `api.addCustomAttribute()`.
- *
- * @param {object} params Function parameters.
- * @param {object} params.tx The current transaction.
- * @param {object} params.msg The outgoing message to decorate.
- */
-function setMetadata({ tx, msg }) {
-  const CONVERSATION_ID = 'llm.conversation_id'
-  const attrs = tx.trace?.custom?.get(DESTINATIONS.TRANS_SCOPE) || {}
-  for (const [key, value] of Object.entries(attrs)) {
-    if (key.startsWith('llm.') && key !== CONVERSATION_ID) {
-      msg[key] = value
-    }
-  }
-}
-
 /**
  * Assigns requestId, conversationId and messageIds for a given
  * chat completion response on the active transaction.
@@ -78,7 +58,7 @@ function assignIdsToTx({ tx, msg, responseId }) {
     tracker.get(responseId) ??
     new LlmTrackedIds({
       requestId: msg.request_id,
-      conversationId: msg.conversation_id
+      conversationId: msg['llm.conversation_id']
     })
   trackedIds.message_ids.push(msg.id)
   tracker.set(responseId, trackedIds)
@@ -115,7 +95,7 @@ function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockR
     completionId: summary.id
   })
   assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg })
-  recordEvent({ agent, type: 'LlmChatCompletionMessage', msg, tx })
+  recordEvent({ agent, type: 'LlmChatCompletionMessage', msg })
 
   bedrockResponse.completions.forEach((content, index) => {
     const chatCompletionMessage = new LlmChatCompletionMessage({
@@ -129,10 +109,10 @@ function recordChatCompletionMessages({ agent, segment, bedrockCommand, bedrockR
       completionId: summary.id
     })
     assignIdsToTx({ tx, responseId: bedrockResponse.requestId, msg: chatCompletionMessage })
-    recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: chatCompletionMessage, tx })
+    recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: chatCompletionMessage })
   })
 
-  recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: summary, tx })
+  recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: summary })
 
   if (err) {
     const llmError = new LlmError({ bedrockResponse, err, summary })
@@ -159,7 +139,7 @@ function recordEmbeddingMessage({ agent, segment, bedrockCommand, bedrockRespons
     isError: err !== null
   })
 
-  recordEvent({ agent, type: 'LlmEmbedding', msg: embedding, tx: agent.getTransaction() })
+  recordEvent({ agent, type: 'LlmEmbedding', msg: embedding })
   if (err) {
     const llmError = new LlmError({ bedrockResponse, err, embedding })
     agent.errors.add(segment.transaction, err, llmError)

tests/unit/llm/chat-completion-message.tap.js

@@ -75,7 +75,7 @@ tap.beforeEach((t) => {
 tap.test('create creates a non-response instance', async (t) => {
   const event = new LlmChatCompletionMessage(t.context)
   t.equal(event.is_response, false)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.completion_id, 'completion-1')
   t.equal(event.sequence, 0)
   t.equal(event.content, 'who are you')
@@ -89,7 +89,7 @@ tap.test('create creates a titan response instance', async (t) => {
   t.context.isResponse = true
   const event = new LlmChatCompletionMessage(t.context)
   t.equal(event.is_response, true)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.completion_id, 'completion-1')
   t.equal(event.sequence, 0)
   t.equal(event.content, 'a response')
@@ -104,7 +104,7 @@ tap.test('create creates a cohere response instance', async (t) => {
   t.context.bedrockResponse.id = 42
   const event = new LlmChatCompletionMessage(t.context)
   t.equal(event.is_response, true)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.completion_id, 'completion-1')
   t.equal(event.sequence, 0)
   t.equal(event.content, 'a response')
@@ -119,7 +119,7 @@ tap.test('create creates a ai21 response instance when response.id is undefined'
   delete t.context.bedrockResponse.id
   const event = new LlmChatCompletionMessage(t.context)
   t.equal(event.is_response, true)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.completion_id, 'completion-1')
   t.equal(event.sequence, 0)
   t.equal(event.content, 'a response')

tests/unit/llm/chat-completion-summary.tap.js

@@ -80,7 +80,7 @@ tap.test('creates a basic summary', async (t) => {
   t.context.bedrockResponse.inputTokenCount = 0
   t.context.bedrockResponse.outputTokenCount = 0
   const event = new LlmChatCompletionSummary(t.context)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.duration, 100)
   t.equal(event['request.max_tokens'], 25)
   t.equal(event['request.temperature'], 0.5)
@@ -94,7 +94,7 @@ tap.test('creates a basic summary', async (t) => {
 tap.test('creates an ai21 summary', async (t) => {
   t.context.bedrockCommand.isAi21 = () => true
   const event = new LlmChatCompletionSummary(t.context)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.duration, 100)
   t.equal(event['request.max_tokens'], 25)
   t.equal(event['request.temperature'], 0.5)
@@ -108,7 +108,7 @@ tap.test('creates an ai21 summary', async (t) => {
 tap.test('creates an claude summary', async (t) => {
   t.context.bedrockCommand.isClaude = () => true
   const event = new LlmChatCompletionSummary(t.context)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.duration, 100)
   t.equal(event['request.max_tokens'], 25)
   t.equal(event['request.temperature'], 0.5)
@@ -122,7 +122,7 @@ tap.test('creates an claude summary', async (t) => {
 tap.test('creates a cohere summary', async (t) => {
   t.context.bedrockCommand.isCohere = () => true
   const event = new LlmChatCompletionSummary(t.context)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.duration, 100)
   t.equal(event['request.max_tokens'], 25)
   t.equal(event['request.temperature'], 0.5)
@@ -136,7 +136,7 @@ tap.test('creates a cohere summary', async (t) => {
 tap.test('creates a llama2 summary', async (t) => {
   t.context.bedrockCommand.isLlama2 = () => true
   const event = new LlmChatCompletionSummary(t.context)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.duration, 100)
   t.equal(event['request.max_tokens'], 25)
   t.equal(event['request.temperature'], 0.5)
@@ -150,7 +150,7 @@ tap.test('creates a llama2 summary', async (t) => {
 tap.test('creates a titan summary', async (t) => {
   t.context.bedrockCommand.isTitan = () => true
   const event = new LlmChatCompletionSummary(t.context)
-  t.equal(event.conversation_id, 'conversation-1')
+  t.equal(event['llm.conversation_id'], 'conversation-1')
   t.equal(event.duration, 100)
   t.equal(event['request.max_tokens'], 25)
   t.equal(event['request.temperature'], 0.5)

tests/unit/llm/event.tap.js

@@ -26,7 +26,8 @@ tap.beforeEach((t) => {
               get(key) {
                 t.equal(key, TRANS_SCOPE)
                 return {
-                  ['llm.conversation_id']: 'conversation-1'
+                  ['llm.conversation_id']: 'conversation-1',
+                  omit: 'me'
                 }
               }
             }
@@ -67,6 +68,8 @@ tap.test('create creates a new instance', async (t) => {
   t.equal(event['response.model'], 'model-1')
   t.equal(event['request.model'], 'model-1')
   t.equal(event['request.max_tokens'], null)
+  t.equal(event['llm.conversation_id'], 'conversation-1')
+  t.equal(event.omit, undefined)
 })
 
 tap.test('serializes the event', (t) => {

tests/versioned/common.js

@@ -130,7 +130,7 @@ function assertChatCompletionSummary({
     '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',
+    'llm.conversation_id': 'convo-id',
     'trace_id': tx.traceId,
     'span_id': tx.trace.root.children[0].id,
     'transaction_id': tx.id,