feat: Add llm attribute to all transactions that contain llm spans fo…

…r openai (#1946)
newrelic · Jan 12, 2024 · 6312629 · 6312629
1 parent cefce24
commit 6312629
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js
@@ -19,6 +19,7 @@ const {
   AI: { OPENAI }
 } = require('../../lib/metrics/names')
 const semver = require('semver')
+const { DESTINATIONS } = require('../config/attribute-filter')
 
 let TRACKING_METRIC = OPENAI.TRACKING_PREFIX
 
@@ -282,6 +283,8 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
               err
             })
           }
+
+          segment.transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
         }
       }
     }
@@ -301,6 +304,8 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
         promise: true,
         // eslint-disable-next-line max-params
         after(_shim, _fn, _name, err, response, segment) {
+          segment.transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
+
           if (!response) {
             // If we get an error, it is possible that `response = null`.
             // In that case, we define it to be an empty object.

diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js
@@ -25,6 +25,7 @@ const fs = require('fs')
 const { version: pkgVersion } = JSON.parse(
   fs.readFileSync(`${__dirname}/node_modules/openai/package.json`)
 )
+const { DESTINATIONS } = require('../../../lib/config/attribute-filter')
 
 tap.test('OpenAI instrumentation - chat completions', (t) => {
   t.autoend()
@@ -343,4 +344,19 @@ tap.test('OpenAI instrumentation - chat completions', (t) => {
       test.end()
     })
   })
+
+  t.test('should add llm attribute to transaction', (test) => {
+    const { client, agent } = t.context
+    helper.runInTransaction(agent, async (tx) => {
+      await client.chat.completions.create({
+        messages: [{ role: 'user', content: 'You are a mathematician.' }]
+      })
+
+      const attributes = tx.trace.attributes.get(DESTINATIONS.TRANS_EVENT)
+      t.equal(attributes.llm, true)
+
+      tx.end()
+      test.end()
+    })
+  })
 })
diff --git a/test/versioned/openai/embeddings.tap.js b/test/versioned/openai/embeddings.tap.js
@@ -23,6 +23,7 @@ const fs = require('fs')
 const { version: pkgVersion } = JSON.parse(
   fs.readFileSync(`${__dirname}/node_modules/openai/package.json`)
 )
+const { DESTINATIONS } = require('../../../lib/config/attribute-filter')
 
 tap.test('OpenAI instrumentation - embedding', (t) => {
   t.autoend()
@@ -153,4 +154,20 @@ tap.test('OpenAI instrumentation - embedding', (t) => {
       test.end()
     })
   })
+
+  t.test('should add llm attribute to transaction', (test) => {
+    const { client, agent } = t.context
+    helper.runInTransaction(agent, async (tx) => {
+      await client.embeddings.create({
+        input: 'This is an embedding test.',
+        model: 'text-embedding-ada-002'
+      })
+
+      const attributes = tx.trace.attributes.get(DESTINATIONS.TRANS_EVENT)
+      t.equal(attributes.llm, true)
+
+      tx.end()
+      test.end()
+    })
+  })
 })