refactor: Updated span names of openai actions to allow AI O11y team …

…to locate relevant spans (#1896)
newrelic · Dec 6, 2023 · b7a644d · b7a644d
1 parent 9973b24
commit b7a644d
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 17 deletions.
diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js
@@ -15,10 +15,12 @@ const LlmTrackedIds = require('../../lib/llm-events/tracked-ids')
 
 const MIN_VERSION = '4.0.0'
 const MIN_STREAM_VERSION = '4.12.2'
-const { AI } = require('../../lib/metrics/names')
+const {
+  AI: { OPENAI }
+} = require('../../lib/metrics/names')
 const semver = require('semver')
 
-let TRACKING_METRIC = AI.TRACKING_PREFIX
+let TRACKING_METRIC = OPENAI.TRACKING_PREFIX
 
 /**
  * Checks if we should skip instrumentation.
@@ -222,7 +224,7 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
   // Update the tracking metric name with the version of the library
   // being instrumented. We do not have access to the version when
   // initially declaring the variable.
-  TRACKING_METRIC = `${TRACKING_METRIC}OpenAI/${shim.pkgVersion}`
+  TRACKING_METRIC = `${TRACKING_METRIC}/${shim.pkgVersion}`
 
   /**
    * Instrumentation is only done to get the response headers and attach
@@ -266,7 +268,7 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
       }
 
       return {
-        name: `${AI.OPEN_AI}/Chat/Completions/Create`,
+        name: OPENAI.COMPLETION,
         promise: true,
         // eslint-disable-next-line max-params
         after(_shim, _fn, _name, err, response, segment) {
@@ -296,7 +298,7 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
     function wrapEmbeddingCreate(shim, embeddingCreate, name, args) {
       const [request] = args
       return {
-        name: `${AI.OPEN_AI}/Embeddings/Create`,
+        name: OPENAI.EMBEDDING,
         promise: true,
         // eslint-disable-next-line max-params
         after(_shim, _fn, _name, err, response, segment) {

diff --git a/lib/metrics/names.js b/lib/metrics/names.js
@@ -166,7 +166,14 @@ const EXPRESS = {
 
 const AI = {
   TRACKING_PREFIX: 'Nodejs/ML/',
-  OPEN_AI: 'AI/OpenAI'
+  EMBEDDING: 'Llm/embedding',
+  COMPLETION: 'Llm/completion'
+}
+
+AI.OPENAI = {
+  TRACKING_PREFIX: `${AI.TRACKING_PREFIX}/OpenAI`,
+  EMBEDDING: `${AI.EMBEDDING}/OpenAI/create`,
+  COMPLETION: `${AI.COMPLETION}/OpenAI/create`
 }
 
 const RESTIFY = {

diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js
@@ -13,7 +13,9 @@
 const tap = require('tap')
 const helper = require('../../lib/agent_helper')
 const { assertSegments } = require('../../lib/metrics_helper')
-const { AI } = require('../../../lib/metrics/names')
+const {
+  AI: { OPENAI }
+} = require('../../../lib/metrics/names')
 const responses = require('./mock-responses')
 const {
   beforeHook,
@@ -55,7 +57,7 @@ tap.test('OpenAI instrumentation - chat completions', (t) => {
       test.doesNotThrow(() => {
         assertSegments(
           tx.trace.root,
-          ['AI/OpenAI/Chat/Completions/Create', [`External/${host}:${port}/chat/completions`]],
+          [OPENAI.COMPLETION, [`External/${host}:${port}/chat/completions`]],
           { exact: false }
         )
       }, 'should have expected segments')
@@ -71,7 +73,7 @@ tap.test('OpenAI instrumentation - chat completions', (t) => {
         messages: [{ role: 'user', content: 'You are a mathematician.' }]
       })
 
-      const metrics = agent.metrics.getOrCreateMetric(`${AI.TRACKING_PREFIX}OpenAI/${pkgVersion}`)
+      const metrics = agent.metrics.getOrCreateMetric(`${OPENAI.TRACKING_PREFIX}/${pkgVersion}`)
       t.equal(metrics.callCount > 0, true)
 
       tx.end()
@@ -138,7 +140,7 @@ tap.test('OpenAI instrumentation - chat completions', (t) => {
         test.doesNotThrow(() => {
           assertSegments(
             tx.trace.root,
-            ['AI/OpenAI/Chat/Completions/Create', [`External/${host}:${port}/chat/completions`]],
+            [OPENAI.COMPLETION, [`External/${host}:${port}/chat/completions`]],
             { exact: false }
           )
         }, 'should have expected segments')

diff --git a/test/versioned/openai/embeddings.tap.js b/test/versioned/openai/embeddings.tap.js
@@ -14,7 +14,9 @@ const tap = require('tap')
 const helper = require('../../lib/agent_helper')
 const { assertSegments } = require('../../lib/metrics_helper')
 const { beforeHook, afterEachHook, afterHook } = require('./common')
-const { AI } = require('../../../lib/metrics/names')
+const {
+  AI: { OPENAI }
+} = require('../../../lib/metrics/names')
 
 const fs = require('fs')
 // have to read and not require because openai does not export the package.json
@@ -44,11 +46,9 @@ tap.test('OpenAI instrumentation - embedding', (t) => {
       test.equal(results.model, 'text-embedding-ada-002-v2')
 
       test.doesNotThrow(() => {
-        assertSegments(
-          tx.trace.root,
-          ['AI/OpenAI/Embeddings/Create', [`External/${host}:${port}/embeddings`]],
-          { exact: false }
-        )
+        assertSegments(tx.trace.root, [OPENAI.EMBEDDING, [`External/${host}:${port}/embeddings`]], {
+          exact: false
+        })
       }, 'should have expected segments')
       tx.end()
       test.end()
@@ -63,7 +63,7 @@ tap.test('OpenAI instrumentation - embedding', (t) => {
         model: 'text-embedding-ada-002'
       })
 
-      const metrics = agent.metrics.getOrCreateMetric(`${AI.TRACKING_PREFIX}OpenAI/${pkgVersion}`)
+      const metrics = agent.metrics.getOrCreateMetric(`${OPENAI.TRACKING_PREFIX}/${pkgVersion}`)
       t.equal(metrics.callCount > 0, true)
 
       tx.end()