feat: Added new API method withLlmCustomAttributes to run a function in a LLM context

api.js

@@ -32,6 +32,7 @@
 const { DESTINATIONS } = require('./lib/config/attribute-filter')
 const parse = require('module-details-from-path')
 const { isSimpleObject } = require('./lib/util/objects')
+const { AsyncLocalStorage } = require('async_hooks')
 
 /*
  *
@@ -1548,7 +1549,7 @@
  * @param {string} params.traceId Identifier for the feedback event.
  * Obtained from {@link getTraceMetadata}.
  * @param {string} params.category A tag for the event.
  * @param {string} params.rating A indicator of how useful the message was.
  * @param {string} [params.message] The message that triggered the event.
  * @param {object} [params.metadata] Additional key-value pairs to associate
  * with the recorded event.
@@ -1902,4 +1903,51 @@
   transaction.ignoreApdex = true
 }
 
+/**
+ * Run a function with the passed in LLM context as the active context and return its return value.
+ *
+ * An example of setting a custom attribute:
+ *
+ *    newrelic.withLlmCustomAttributes({'llm.someAttribute': 'someValue'}, () => {
+ *      return;
+ *    })
+ * @param {Object} context LLM custom attributes context
+ * @param {Function} callback The function to execute in context.
+ */
+API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context, callback) {
+  const metric = this.agent.metrics.getOrCreateMetric(
+    NAMES.SUPPORTABILITY.API + '/withLlmCustomAttributes'
+  )
+  metric.incrementCallCount()
+
+  const transaction = this.agent.tracer.getTransaction()
+
+  if (!callback || typeof callback !== 'function') {
+    logger.warn('withLlmCustomAttributes must be used with a valid callback')
+    return
+  }
+
+  if (!transaction) {
+    logger.warn('withLlmCustomAttributes must be called within the scope of a transaction.')
+    return callback()
+  }
+
+  for (const [key, value] of Object.entries(context)) {
+    if (typeof value === 'object' || typeof value === 'function') {
+      logger.warn(`Invalid attribute type for ${key}. Skipped.`)
+      delete context[key]
+    } else if (key.indexOf('llm.') !== 0) {
+      logger.warn(`Invalid attribute name ${key}. Renamed to "llm.${key}".`)
+      delete context[key]
+      context[`llm.${key}`] = value
+    }
+  }
+
+  transaction._llmContextManager = transaction._llmContextManager || new AsyncLocalStorage()
+  const parentContext = transaction._llmContextManager.getStore()
+
+  const fullContext = Object.assign({}, parentContext || {}, context || {})
+  return transaction._llmContextManager.run(fullContext, callback)
+}
+
 module.exports = API

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

@@ -18,6 +18,7 @@ const { DESTINATIONS } = require('../../../config/attribute-filter')
 const { AI } = require('../../../metrics/names')
 const { RecorderSpec } = require('../../../shim/specs')
 const InstrumentationDescriptor = require('../../../instrumentation-descriptor')
+const { extractLlmContext } = require('../../../util/llm-utils')
 
 let TRACKING_METRIC
 
@@ -55,7 +56,12 @@ function isStreamingEnabled({ commandName, config }) {
  */
 function recordEvent({ agent, type, msg }) {
   msg.serialize()
-  agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
+  const llmContext = extractLlmContext(agent)
+
+  agent.customEventAggregator.add([
+    { type, timestamp: Date.now() },
+    Object.assign({}, msg, llmContext)
+  ])
 }
 
 /**

lib/instrumentation/langchain/common.js

@@ -7,6 +7,7 @@
 const {
   AI: { LANGCHAIN }
 } = require('../../metrics/names')
+const { extractLlmContext } = require('../../util/llm-utils')
 
 const common = module.exports
 
@@ -49,7 +50,12 @@ common.mergeMetadata = function mergeMetadata(localMeta = {}, paramsMeta = {}) {
  */
 common.recordEvent = function recordEvent({ agent, type, msg, pkgVersion }) {
   agent.metrics.getOrCreateMetric(`${LANGCHAIN.TRACKING_PREFIX}/${pkgVersion}`).incrementCallCount()
-  agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
+  const llmContext = extractLlmContext(agent)
+
+  agent.customEventAggregator.add([
+    { type, timestamp: Date.now() },
+    Object.assign({}, msg, llmContext)
+  ])
 }
 
 /**

lib/instrumentation/openai.js

@@ -12,6 +12,7 @@ const {
   LlmErrorMessage
 } = require('../../lib/llm-events/openai')
 const { RecorderSpec } = require('../../lib/shim/specs')
+const { extractLlmContext } = require('../util/llm-utils')
 
 const MIN_VERSION = '4.0.0'
 const MIN_STREAM_VERSION = '4.12.2'
@@ -75,7 +76,12 @@ function decorateSegment({ shim, result, apiKey }) {
  * @param {object} params.msg LLM event
  */
 function recordEvent({ agent, type, msg }) {
-  agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
+  const llmContext = extractLlmContext(agent)
+
+  agent.customEventAggregator.add([
+    { type, timestamp: Date.now() },
+    Object.assign({}, msg, llmContext)
+  ])
 }
 
 /**

lib/util/llm-utils.js

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2020 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+exports = module.exports = { extractLlmContext, extractLlmAttribtues }
+
+/**
+ * Extract LLM attributes from the LLM context
+ *
+ * @param {Object} context LLM context object
+ * @returns {Object} LLM custom attributes
+ */
+function extractLlmAttribtues(context) {
+  return Object.keys(context || {}).reduce((result, key) => {
+    if (key.indexOf('llm.') === 0) {
+      result[key] = context[key]
+    }
+    return result
+  }, {})
+}
+
+/**
+ * Extract LLM context from the active transaction
+ *
+ * @param {Agent} agent NR agent instance
+ * @returns {Object} LLM context object
+ */
+function extractLlmContext(agent) {
+  const context = agent.tracer.getTransaction()._llmContextManager?.getStore() || {}
+  return extractLlmAttribtues(context)
+}

test/unit/api/api-llm.test.js

@@ -121,6 +121,49 @@ tap.test('Agent API LLM methods', (t) => {
     })
   })
 
+  t.test('withLlmCustomAttributes', (t) => {
+    const { api } = t.context
+    helper.runInTransaction(api.agent, (tx) => {
+      t.context.agent.tracer.getTransaction = () => {
+        return tx
+      }
+
+      t.doesNotThrow(() => {
+        api.withLlmCustomAttributes(null, null)
+        t.equal(loggerMock.warn.callCount, 1)
+      })
+
+      api.withLlmCustomAttributes(
+        {
+          'toRename': 'value1',
+          'llm.number': 1,
+          'llm.boolean': true,
+          'toDelete': () => {},
+          'toDelete2': {},
+          'toDelete3': []
+        },
+        () => {
+          const contextManager = tx._llmContextManager
+          const parentContext = contextManager.getStore()
+          t.equal(parentContext['llm.toRename'], 'value1')
+          t.notOk(parentContext.toDelete)
+          t.notOk(parentContext.toDelete2)
+          t.notOk(parentContext.toDelete3)
+          t.equal(parentContext['llm.number'], 1)
+          t.equal(parentContext['llm.boolean'], true)
+
+          api.withLlmCustomAttributes({ 'llm.someAttribute': 'someValue' }, () => {
+            const contextManager = tx._llmContextManager
+            const context = contextManager.getStore()
+            t.equal(context[`llm.toRename`], 'value1')
+            t.equal(context['llm.someAttribute'], 'someValue')
+            t.end()
+          })
+        }
+      )
+    })
+  })
+
   t.test('setLlmTokenCount should register callback to calculate token counts', async (t) => {
     const { api, agent } = t.context
     function callback(model, content) {

test/unit/api/stub.test.js

@@ -8,7 +8,7 @@
 const tap = require('tap')
 const API = require('../../../stub_api')
 
-const EXPECTED_API_COUNT = 36
+const EXPECTED_API_COUNT = 37
 
 tap.test('Agent API - Stubbed Agent API', (t) => {
   t.autoend()

test/unit/instrumentation/openai.test.js

@@ -119,5 +119,6 @@ test('openai unit tests', (t) => {
     t.equal(isWrapped, false, 'should not wrap chat completions create')
     t.end()
   })
+
   t.end()
 })

test/unit/util/llm-utils.test.js

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2023 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const tap = require('tap')
+const { extractLlmAttribtues, extractLlmContext } = require('../../../lib/util/llm-utils')
+const { AsyncLocalStorage } = require('async_hooks')
+
+tap.test('extractLlmAttributes', (t) => {
+  const context = {
+    'skip': 1,
+    'llm.get': 2,
+    'fllm.skip': 3
+  }
+
+  const llmContext = extractLlmAttribtues(context)
+  t.notOk(llmContext.skip)
+  t.notOk(llmContext['fllm.skip'])
+  t.equal(llmContext['llm.get'], 2)
+  t.end()
+})
+
+tap.test('extractLlmContext', (t) => {
+  const tx = {
+    _llmContextManager: new AsyncLocalStorage()
+  }
+  const agent = {
+    tracer: {
+      getTransaction: () => {
+        return tx
+      }
+    }
+  }
+
+  tx._llmContextManager.run(null, () => {
+    const llmContext = extractLlmContext(agent)
+    t.equal(typeof llmContext, 'object')
+    t.equal(Object.entries(llmContext).length, 0)
+  })
+
+  tx._llmContextManager.run({ 'llm.test': 1, 'skip': 2 }, () => {
+    const llmContext = extractLlmContext(agent)
+    t.equal(llmContext['llm.test'], 1)
+    t.notOk(llmContext.skip)
+    t.end()
+  })
+})

test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js

@@ -158,6 +158,32 @@ tap.afterEach(async (t) => {
     }
   )
 
+  tap.test(
+    `${modelId}:  supports custom attributes on LlmChatCompletionMessage(s) and LlmChatCompletionSummary events`,
+    (t) => {
+      const { bedrock, client, agent } = t.context
+      const prompt = `text ${resKey} ultimate question`
+      const input = requests[resKey](prompt, modelId)
+      const command = new bedrock.InvokeModelCommand(input)
+
+      const api = helper.getAgentApi()
+      helper.runInTransaction(agent, async (tx) => {
+        api.addCustomAttribute('llm.conversation_id', 'convo-id')
+        api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => {
+          await client.send(command)
+          const events = agent.customEventAggregator.events.toArray()
+
+          const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0]
+          const [, message] = chatSummary
+          t.equal(message['llm.contextAttribute'], 'someValue')
+
+          tx.end()
+          t.end()
+        })
+      })
+    }
+  )
+
   tap.test(`${modelId}: text answer (streamed)`, (t) => {
     if (modelId.includes('ai21')) {
       t.skip('model does not support streaming')

test/versioned/langchain/runnables.tap.js

@@ -52,7 +52,6 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => {
 
   t.test('should create langchain events for every invoke call', (t) => {
     const { agent, prompt, outputParser, model } = t.context
-
     helper.runInTransaction(agent, async (tx) => {
       const input = { topic: 'scientist' }
       const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] }
@@ -95,11 +94,31 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => {
     })
   })
 
+  t.test('should support custom attributes on the LLM events', (t) => {
+    const { agent, prompt, outputParser, model } = t.context
+    const api = helper.getAgentApi()
+    helper.runInTransaction(agent, async (tx) => {
+      api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => {
+        const input = { topic: 'scientist' }
+        const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] }
+
+        const chain = prompt.pipe(model).pipe(outputParser)
+        await chain.invoke(input, options)
+        const events = agent.customEventAggregator.events.toArray()
+
+        const [[, message]] = events
+        t.equal(message['llm.contextAttribute'], 'someValue')
+
+        tx.end()
+        t.end()
+      })
+    })
+  })
+
   t.test(
     'should create langchain events for every invoke call on chat prompt + model + parser',
     (t) => {
       const { agent, prompt, outputParser, model } = t.context
-
       helper.runInTransaction(agent, async (tx) => {
         const input = { topic: 'scientist' }
         const options = { metadata: { key: 'value', hello: 'world' }, tags: ['tag1', 'tag2'] }