From d11d100a37ad04201a9065363b1bc6b2d6fe36f2 Mon Sep 17 00:00:00 2001
From: James Sumners <jsumners@newrelic.com>
Date: Tue, 28 Nov 2023 10:59:51 -0500
Subject: [PATCH] feat: Add version tracking metric on OpenAI events (#1886)

---
 lib/instrumentation/openai.js                 |  8 +++++++
 lib/metrics/names.js                          |  1 +
 test/versioned/openai/chat-completions.tap.js | 16 +++++++++++++
 test/versioned/openai/embeddings.tap.js       | 23 +++++++++++++++++++
 4 files changed, 48 insertions(+)

diff --git a/lib/instrumentation/openai.js b/lib/instrumentation/openai.js
index 877d413cc9..24162af182 100644
--- a/lib/instrumentation/openai.js
+++ b/lib/instrumentation/openai.js
@@ -18,6 +18,8 @@ const MIN_STREAM_VERSION = '4.12.2'
 const { AI } = require('../../lib/metrics/names')
 const semver = require('semver')
 
+let TRACKING_METRIC = AI.TRACKING_PREFIX
+
 /**
  * Checks if we should skip instrumentation.
  * Currently it checks if `ai_monitoring.enabled` is true
@@ -46,6 +48,11 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
     return
   }
 
+  // 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}`
+
   /**
    * Adds apiKey and response headers to the active segment
    * on symbols
@@ -76,6 +83,7 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
    * @param {object} msg LLM event
    */
   function recordEvent(type, msg) {
+    agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount()
     msg = agent?.llm?.metadata ? { ...agent.llm.metadata, ...msg } : msg
     agent.customEventAggregator.add([{ type, timestamp: Date.now() }, msg])
   }
diff --git a/lib/metrics/names.js b/lib/metrics/names.js
index 7cef2f6c8c..23ead7925b 100644
--- a/lib/metrics/names.js
+++ b/lib/metrics/names.js
@@ -165,6 +165,7 @@ const EXPRESS = {
 }
 
 const AI = {
+  TRACKING_PREFIX: 'Nodejs/ML/',
   OPEN_AI: 'AI/OpenAI'
 }
 
diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js
index bc09af31dd..449d598352 100644
--- a/test/versioned/openai/chat-completions.tap.js
+++ b/test/versioned/openai/chat-completions.tap.js
@@ -13,6 +13,7 @@
 const tap = require('tap')
 const helper = require('../../lib/agent_helper')
 const { assertSegments } = require('../../lib/metrics_helper')
+const { AI } = require('../../../lib/metrics/names')
 const responses = require('./mock-responses')
 const {
   beforeHook,
@@ -63,6 +64,21 @@ tap.test('OpenAI instrumentation - chat completions', (t) => {
     })
   })
 
+  t.test('should increment tracking metric for each chat completion event', (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 metrics = agent.metrics.getOrCreateMetric(`${AI.TRACKING_PREFIX}OpenAI/${pkgVersion}`)
+      t.equal(metrics.callCount > 0, true)
+
+      tx.end()
+      test.end()
+    })
+  })
+
   t.test('should create chat completion message and summary for every message sent', (test) => {
     const { client, agent } = t.context
     helper.runInTransaction(agent, async (tx) => {
diff --git a/test/versioned/openai/embeddings.tap.js b/test/versioned/openai/embeddings.tap.js
index 1da5d1a89c..712d1f5f08 100644
--- a/test/versioned/openai/embeddings.tap.js
+++ b/test/versioned/openai/embeddings.tap.js
@@ -14,6 +14,13 @@ 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 fs = require('fs')
+// have to read and not require because openai does not export the package.json
+const { version: pkgVersion } = JSON.parse(
+  fs.readFileSync(`${__dirname}/node_modules/openai/package.json`)
+)
 
 tap.test('OpenAI instrumentation - embedding', (t) => {
   t.autoend()
@@ -48,6 +55,22 @@ tap.test('OpenAI instrumentation - embedding', (t) => {
     })
   })
 
+  t.test('should increment tracking metric for each embedding event', (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 metrics = agent.metrics.getOrCreateMetric(`${AI.TRACKING_PREFIX}OpenAI/${pkgVersion}`)
+      t.equal(metrics.callCount > 0, true)
+
+      tx.end()
+      test.end()
+    })
+  })
+
   t.test('should create an embedding message', (test) => {
     const { client, agent } = t.context
     helper.runInTransaction(agent, async (tx) => {