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

api.js

@@ -1932,7 +1932,7 @@ API.prototype.withLlmCustomAttributes = function withLlmCustomAttributes(context
     return callback()
   }
 
-  for (const [key, value] of Object.entries(context)) {
+  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]

lib/util/llm-utils.js

@@ -5,16 +5,16 @@
 
 'use strict'
 
-exports = module.exports = { extractLlmContext, extractLlmAttribtues }
+exports = module.exports = { extractLlmContext, extractLlmAttributes }
 
 /**
  * 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) => {
+function extractLlmAttributes(context) {
+  return Object.keys(context).reduce((result, key) => {
     if (key.indexOf('llm.') === 0) {
       result[key] = context[key]
     }
@@ -29,6 +29,6 @@ function extractLlmAttribtues(context) {
  * @returns {Object} LLM context object
  */
 function extractLlmContext(agent) {
-  const context = agent.tracer.getTransaction()._llmContextManager?.getStore() || {}
-  return extractLlmAttribtues(context)
+  const context = agent.tracer.getTransaction()?._llmContextManager?.getStore() || {}
+  return extractLlmAttributes(context)
 }

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

@@ -121,27 +121,58 @@ tap.test('Agent API LLM methods', (t) => {
     })
   })
 
-  t.test('withLlmCustomAttributes', (t) => {
+  t.test('withLlmCustomAttributes should handle no active transaction', (t) => {
     const { api } = t.context
-    t.doesNotThrow(() => {
+    t.autoend()
+
+    t.equal(
+      api.withLlmCustomAttributes({ test: 1 }, () => {
+        t.equal(loggerMock.warn.callCount, 1)
+        return 1
+      }),
+      1
+    )
+  })
+
+  t.test('withLlmCustomAttributes should handle an empty store', (t) => {
+    const { api } = t.context
+    const agent = api.agent
+
+    t.autoend()
+    helper.runInTransaction(api.agent, (tx) => {
+      agent.tracer.getTransaction = () => {
+        return tx
+      }
       t.equal(
-        api.withLlmCustomAttributes({ test: 1 }, () => {
-          t.equal(loggerMock.warn.callCount, 1)
+        api.withLlmCustomAttributes(null, () => {
           return 1
         }),
         1
       )
     })
+  })
+
+  t.test('withLlmCustomAttributes should handle no callback', (t) => {
+    const { api } = t.context
+    const agent = api.agent
+    t.autoend()
     helper.runInTransaction(api.agent, (tx) => {
-      t.context.agent.tracer.getTransaction = () => {
+      agent.tracer.getTransaction = () => {
         return tx
       }
+      api.withLlmCustomAttributes({ test: 1 }, null)
+      t.equal(loggerMock.warn.callCount, 1)
+    })
+  })
 
-      t.doesNotThrow(() => {
-        api.withLlmCustomAttributes(null, null)
-        t.equal(loggerMock.warn.callCount, 2)
-      })
-
+  t.test('withLlmCustomAttributes should normalize attributes', (t) => {
+    const { api } = t.context
+    const agent = api.agent
+    t.autoend()
+    helper.runInTransaction(api.agent, (tx) => {
+      agent.tracer.getTransaction = () => {
+        return tx
+      }
       api.withLlmCustomAttributes(
         {
           'toRename': 'value1',
@@ -160,19 +191,40 @@ tap.test('Agent API LLM methods', (t) => {
           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('withLlmCustomAttributes should support branching', (t) => {
+    const { api } = t.context
+    const agent = api.agent
+    t.autoend()
+    helper.runInTransaction(api.agent, (tx) => {
+      agent.tracer.getTransaction = () => {
+        return tx
+      }
+      api.withLlmCustomAttributes({ 'llm.step': '1', 'llm.path': 'root' }, () => {
+        const contextManager = tx._llmContextManager
+        const context = contextManager.getStore()
+        t.equal(context[`llm.step`], '1')
+        t.equal(context['llm.path'], 'root')
+        api.withLlmCustomAttributes({ 'llm.step': '1.1', 'llm.path': 'root/1' }, () => {
+          const contextManager = tx._llmContextManager
+          const context = contextManager.getStore()
+          t.equal(context[`llm.step`], '1.1')
+          t.equal(context['llm.path'], 'root/1')
+        })
+        api.withLlmCustomAttributes({ 'llm.step': '1.2', 'llm.path': 'root/2' }, () => {
+          const contextManager = tx._llmContextManager
+          const context = contextManager.getStore()
+          t.equal(context[`llm.step`], '1.2')
+          t.equal(context['llm.path'], 'root/2')
+        })
+      })
+    })
+  })
+
   t.test('setLlmTokenCount should register callback to calculate token counts', async (t) => {
     const { api, agent } = t.context
     function callback(model, content) {

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

@@ -6,7 +6,7 @@
 'use strict'
 
 const tap = require('tap')
-const { extractLlmAttribtues, extractLlmContext } = require('../../../lib/util/llm-utils')
+const { extractLlmAttributes, extractLlmContext } = require('../../../lib/util/llm-utils')
 const { AsyncLocalStorage } = require('async_hooks')
 
 tap.test('extractLlmAttributes', (t) => {
@@ -16,35 +16,57 @@ tap.test('extractLlmAttributes', (t) => {
     'fllm.skip': 3
   }
 
-  const llmContext = extractLlmAttribtues(context)
+  const llmContext = extractLlmAttributes(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
+  let tx
+  let agent
+  t.autoend()
+  t.beforeEach(() => {
+    tx = {
+      _llmContextManager: new AsyncLocalStorage()
+    }
+    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)
+  t.test('handle empty context', (t) => {
+    t.autoend()
+    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()
+  t.test('extract LLM context', (t) => {
+    t.autoend()
+    tx._llmContextManager.run({ 'llm.test': 1, 'skip': 2 }, () => {
+      const llmContext = extractLlmContext(agent)
+      t.equal(llmContext['llm.test'], 1)
+      t.notOk(llmContext.skip)
+    })
+  })
+
+  t.test('no transaction', (t) => {
+    t.autoend()
+    agent.tracer.getTransaction = () => {
+      return null
+    }
+    tx._llmContextManager.run(null, () => {
+      const llmContext = extractLlmContext(agent)
+      t.equal(typeof llmContext, 'object')
+      t.equal(Object.entries(llmContext).length, 0)
+    })
   })
 })