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

api.js

@@ -1548,7 +1548,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 +1902,47 @@
   transaction.ignoreApdex = true
 }
 
+/**
+ * Runs a function synchronously within a provided LLM custom attributes context and returns its return value.
+ *
+ * An example of setting a custom attribute:
+ *
+ *    newrelic.withLlmCustomAttributes({'llm.someAttribute': 'someVallue'}, () => {
+ *      return;
+ *    })
+ * @param {Object} context LLM custom attributes context
+ * @param {Function} callback synchronous function called within the 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 (!transaction) {
+    logger.warn('withLlmCustomAttributes must be called within the scope of a transaction.')
+    return callback?.()
+  }
+
+  for (const key in context) {
+    if (Object.hasOwn(context, key)) {
+      const value = context[key]
+      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
+      }
+    }
+  }
+
+  const parentContext = this.agent._contextManager.getContext()
+
+  const fullContext = parentContext ? Object.assign(parentContext, context || {}) : context
+  return this.agent._contextManager.runInContext(fullContext, callback, this)
+}
+
 module.exports = API

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

@@ -55,7 +55,18 @@ function isStreamingEnabled({ commandName, config }) {
  */
 function recordEvent({ agent, type, msg }) {
   msg.serialize()
-  agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
+  const context = agent._contextManager.getContext()
+  const llmContext = Object.keys(context || {}).reduce((result, key) => {
+    if (key.indexOf('llm.') === 0) {
+      result[key] = context[key]
+    }
+    return result
+  }, {})
+
+  agent.customEventAggregator.add([
+    { type, timestamp: Date.now() },
+    Object.assign({}, msg, llmContext || {})
+  ])
 }
 
 /**

lib/instrumentation/langchain/common.js

@@ -49,7 +49,18 @@ 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 context = agent._contextManager.getContext()
+  const llmContext = Object.keys(context || {}).reduce((result, key) => {
+    if (key.indexOf('llm.') === 0) {
+      result[key] = context[key]
+    }
+    return result
+  }, {})
+
+  agent.customEventAggregator.add([
+    { type, timestamp: Date.now() },
+    Object.assign({}, msg, llmContext || {})
+  ])
 }
 
 /**

lib/instrumentation/openai.js

@@ -75,7 +75,18 @@ 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 context = agent._contextManager.getContext()
+  const llmContext = Object.keys(context || {}).reduce((result, key) => {
+    if (key.indexOf('llm.') === 0) {
+      result[key] = context[key]
+    }
+    return result
+  }, {})
+
+  agent.customEventAggregator.add([
+    { type, timestamp: Date.now() },
+    Object.assign({}, msg, llmContext || {})
+  ])
 }
 
 /**

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

@@ -121,6 +121,48 @@ tap.test('Agent API LLM methods', (t) => {
     })
   })
 
+  t.test('withLlmCustomAttributes', (t) => {
+    const { api } = t.context
+    helper.runInTransaction(api.agent, (tx) => {
+      const contextManager = api.agent._contextManager
+      t.context.agent.tracer.getTransaction = () => {
+        return tx
+      }
+      contextManager.setContext({
+        initial: true
+      })
+
+      api.withLlmCustomAttributes(
+        {
+          'toRename': 'value1',
+          'llm.number': 1,
+          'llm.boolean': true,
+          'toDelete': () => {},
+          'toDelete2': {},
+          'toDelete3': []
+        },
+        () => {
+          const parentContext = contextManager.getContext()
+          t.ok(parentContext.initial)
+          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 context = contextManager.getContext()
+            t.ok(context.initial)
+            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,27 @@ test('openai unit tests', (t) => {
     t.equal(isWrapped, false, 'should not wrap chat completions create')
     t.end()
   })
+
+  t.test('should record LLM custom events with attributes', (t) => {
+    const { shim, agent, initialize } = t.context
+    shim.pkgVersion = '4.12.2'
+    const MockOpenAi = getMockModule()
+    agent.config.ai_monitoring.record_content = { enabled: true }
+    initialize(agent, MockOpenAi, 'openai', shim)
+    const completions = new MockOpenAi.Chat.Completions()
+    agent._contextManager.setContext({ initial: true })
+    const api = helper.getAgentApi()
+    helper.runInTransaction(agent, () => {
+      api.withLlmCustomAttributes({ 'llm.attribute': `someValue` }, async () => {
+        await completions.create({ stream: false, messages: [{ role: 'user', content: 'Hello' }] })
+        const events = agent.customEventAggregator.events.toArray()
+        const [[, message]] = events
+        t.notOk(message.initial)
+        t.equal(message['llm.attribute'], 'someValue')
+        t.end()
+      })
+    })
+  })
+
   t.end()
 })

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

@@ -135,25 +135,30 @@ tap.afterEach(async (t) => {
       const api = helper.getAgentApi()
       helper.runInTransaction(agent, async (tx) => {
         api.addCustomAttribute('llm.conversation_id', 'convo-id')
-        await client.send(command)
-        const events = agent.customEventAggregator.events.toArray()
-        t.equal(events.length, 3)
-        const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0]
-        const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage')
-
-        t.llmMessages({
-          modelId,
-          prompt,
-          resContent: '42',
-          tx,
-          expectedId: modelId.includes('ai21') || modelId.includes('cohere') ? '1234' : null,
-          chatMsgs
+        api.withLlmCustomAttributes({ 'llm.contextAttribute': 'someValue' }, async () => {
+          await client.send(command)
+          const events = agent.customEventAggregator.events.toArray()
+          t.equal(events.length, 3)
+          const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0]
+          const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage')
+
+          t.llmMessages({
+            modelId,
+            prompt,
+            resContent: '42',
+            tx,
+            expectedId: modelId.includes('ai21') || modelId.includes('cohere') ? '1234' : null,
+            chatMsgs
+          })
+
+          t.llmSummary({ tx, modelId, chatSummary })
+
+          const [, message] = chatSummary
+          t.ok(message['llm.contextAttribute'])
+
+          tx.end()
+          t.end()
         })
-
-        t.llmSummary({ tx, modelId, chatSummary })
-
-        tx.end()
-        t.end()
       })
     }
   )

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'] }
@@ -99,39 +98,45 @@ tap.test('Langchain instrumentation - runnable sequence', (t) => {
     'should create langchain events for every invoke call on chat prompt + model + parser',
     (t) => {
       const { agent, prompt, outputParser, model } = t.context
-
+      const api = helper.getAgentApi()
+      agent._contextManager.setContext({ initial: true })
       helper.runInTransaction(agent, async (tx) => {
-        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 langchainEvents = filterLangchainEvents(events)
-        const langChainMessageEvents = filterLangchainEventsByType(
-          langchainEvents,
-          'LlmChatCompletionMessage'
-        )
-        const langChainSummaryEvents = filterLangchainEventsByType(
-          langchainEvents,
-          'LlmChatCompletionSummary'
-        )
-
-        t.langchainSummary({
-          tx,
-          chatSummary: langChainSummaryEvents[0]
-        })
-
-        t.langchainMessages({
-          tx,
-          chatMsgs: langChainMessageEvents,
-          chatSummary: langChainSummaryEvents[0][1]
+        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 langchainEvents = filterLangchainEvents(events)
+          const langChainMessageEvents = filterLangchainEventsByType(
+            langchainEvents,
+            'LlmChatCompletionMessage'
+          )
+          const langChainSummaryEvents = filterLangchainEventsByType(
+            langchainEvents,
+            'LlmChatCompletionSummary'
+          )
+
+          t.langchainSummary({
+            tx,
+            chatSummary: langChainSummaryEvents[0]
+          })
+          const [[, message]] = events
+
+          t.ok(message['llm.contextAttribute'])
+
+          t.langchainMessages({
+            tx,
+            chatMsgs: langChainMessageEvents,
+            chatSummary: langChainSummaryEvents[0][1]
+          })
+
+          tx.end()
+          t.end()
         })
-
-        tx.end()
-        t.end()
       })
     }
   )