feat: Add version tracking metric on OpenAI events

lib/instrumentation/openai.js

@@ -13,9 +13,12 @@ const {
 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 semver = require('semver')
 
+let TRACKING_METRIC = AI.TRACKING_PREFIX
+
 /**
  * Checks if we should skip instrumentation.
  * Currently it checks if `ai_monitoring.enabled` is true
@@ -31,10 +34,11 @@ function shouldSkipInstrumentation(config, shim) {
     return true
   }
 
-  const { version: pkgVersion } = shim.require('./package.json')
+  const { pkgVersion } = shim
   return semver.lt(pkgVersion, MIN_VERSION)
 }
 
+// eslint-disable-next-line sonarjs/cognitive-complexity
 module.exports = function initialize(agent, openai, moduleName, shim) {
   if (shouldSkipInstrumentation(agent.config, shim)) {
     shim.logger.debug(
@@ -43,6 +47,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}${shim.pkgVersion}`
+
   /**
    * Adds apiKey and response headers to the active segment
    * on symbols
@@ -53,11 +62,13 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
   function decorateSegment(result, apiKey) {
     const segment = shim.getActiveSegment()
 
-    if (segment) {
-      segment[openAiApiKey] = apiKey
-      segment[openAiHeaders] =
-        result?.response?.headers && Object.fromEntries(result.response.headers)
+    if (!segment) {
+      return
     }
+
+    segment[openAiApiKey] = apiKey
+    segment[openAiHeaders] =
+      result?.response?.headers && Object.fromEntries(result.response.headers)
   }
 
   /**
@@ -67,6 +78,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])
   }
@@ -129,7 +141,12 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
    * @param {object} params.response chat completion response
    */
   function recordChatCompletionMessages({ segment, request, response }) {
+    response.headers = segment[openAiHeaders]
+    response.api_key = segment[openAiApiKey]
     const tx = segment.transaction
+    // explicitly end segment to consistent duration
+    // for both LLM events and the segment
+    segment.end()
     const completionSummary = new LlmChatCompletionSummary({
       agent,
       segment,
@@ -155,6 +172,60 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
     })
 
     recordEvent('LlmChatCompletionSummary', completionSummary)
+    delete response.headers
+    delete response.api_key
+  }
+
+  /*
+   * Chat completions create can return a stream once promise resolves
+   * This wraps the iterator which is a generator function
+   * We will call the original iterator, intercept chunks and yield
+   * to the original. On complete we will construct the new message object
+   * with what we have seen in the stream and create the chat completion
+   * messages
+   *
+   */
+  function instrumentStream({ request, response, segment }) {
+    if (semver.lt(shim.pkgVersion, MIN_STREAM_VERSION)) {
+      shim.logger.warn(
+        `instrumenting chat completion streams is only supported with openai version ${MIN_STREAM_VERSION}+`
+      )
+      return
+    }
+
+    shim.wrap(response, 'iterator', function wrapIterator(shim, orig) {
+      return async function* wrappedIterator() {
+        let content = ''
+        let role = ''
+        let chunk
+        try {
+          const iterator = orig.apply(this, arguments)
+
+          for await (chunk of iterator) {
+            if (chunk.choices[0]?.delta?.role) {
+              role = chunk.choices[0]?.delta?.role
+            }
+
+            content += chunk.choices[0]?.delta?.content ?? ''
+            yield chunk
+          }
+        } catch (err) {
+          // TODO: I cannot figure out how to get an error to throw
+          // this may be unused code
+        } finally {
+          chunk.choices[0].message = { role, content }
+          // update segment duration since we want to extend the time it took to
+          // handle the stream
+          segment.touch()
+
+          recordChatCompletionMessages({
+            segment,
+            request,
+            response: chunk
+          })
+        }
+      }
+    })
   }
 
   /**
@@ -173,22 +244,15 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
         promise: true,
         // eslint-disable-next-line max-params
         after(_shim, _fn, _name, err, response, segment) {
-          response.headers = segment[openAiHeaders]
-          response.api_key = segment[openAiApiKey]
-
-          // TODO: add LlmErrorMessage on failure
-          // and exit
-          // See: https://github.com/newrelic/node-newrelic/issues/1845
-          // if (err) {}
-
-          recordChatCompletionMessages({
-            segment,
-            request,
-            response
-          })
-          // cleanup keys on response before returning to user code
-          delete response.api_key
-          delete response.headers
+          if (request.stream) {
+            instrumentStream({ request, response, segment })
+          } else {
+            recordChatCompletionMessages({
+              segment,
+              request,
+              response
+            })
+          }
         }
       }
     }
@@ -210,6 +274,9 @@ module.exports = function initialize(agent, openai, moduleName, shim) {
         after(_shim, _fn, _name, err, response, segment) {
           response.headers = segment[openAiHeaders]
           response.api_key = segment[openAiApiKey]
+          // explicitly end segment to consistent duration
+          // for both LLM events and the segment
+          segment.end()
           const embedding = new LlmEmbedding({
             agent,
             segment,

lib/llm-events/openai/event.js

@@ -27,7 +27,7 @@ module.exports = class LlmEvent {
      */
     if (responseAttrs) {
       this['request.model'] = request.model || request.engine
-      this.duration = segment?.getExclusiveDurationInMillis()
+      this.duration = segment?.getDurationInMillis()
       this.api_key_last_four_digits = response?.api_key && `sk-${response.api_key.slice(-4)}`
       this.responseAttrs(response)
     }

lib/metrics/names.js

@@ -165,6 +165,7 @@ const EXPRESS = {
 }
 
 const AI = {
+  TRACKING_PREFIX: 'Nodejs/ML/OpenAI/',
   OPEN_AI: 'AI/OpenAI'
 }
 

lib/shim/index.js

@@ -32,13 +32,13 @@
  * @param resolvedName
  * @param shimName
  */
-function createShimFromType(type, agent, moduleName, resolvedName, shimName) {
+function createShimFromType(type, agent, moduleName, resolvedName, shimName, pkgVersion) {
   let shim = null
   if (properties.hasOwn(SHIM_TYPE_MAP, type)) {
     const ShimClass = SHIM_TYPE_MAP[type]
-    shim = new ShimClass(agent, moduleName, resolvedName, shimName)
+    shim = new ShimClass(agent, moduleName, resolvedName, shimName, pkgVersion)
   } else {
-    shim = new Shim(agent, moduleName, resolvedName, shimName)
+    shim = new Shim(agent, moduleName, resolvedName, shimName, pkgVersion)
   }
   return shim
 }

lib/shim/shim.js

@@ -30,9 +30,10 @@ const fnApply = Function.prototype.apply
  * @param {string}  moduleName    - The name of the module being instrumented.
  * @param {string}  resolvedName  - The full path to the loaded module.
  * @param {string}  shimName      - Used to persist shim ids across different instances. This is
+ * @param {string} pkgVersion     - version of package getting instrumented
  * applicable to instrument that compliments each other across libraries(i.e - koa + koa-route/koa-router)
  */
-function Shim(agent, moduleName, resolvedName, shimName) {
+function Shim(agent, moduleName, resolvedName, shimName, pkgVersion) {
   if (!agent || !moduleName) {
     throw new Error('Shim must be initialized with an agent and module name.')
   }
@@ -44,6 +45,7 @@ function Shim(agent, moduleName, resolvedName, shimName) {
   this._debug = false
   this.defineProperty(this, 'moduleName', moduleName)
   this.assignId(shimName)
+  this.pkgVersion = pkgVersion
 
   // Determine the root directory of the module.
   let moduleRoot = null

lib/shimmer.js

@@ -536,7 +536,8 @@ function instrumentPostLoad(agent, nodule, moduleName, resolvedName, esmResolver
       agent,
       moduleName,
       resolvedName,
-      instrumentation.shimName
+      instrumentation.shimName,
+      pkgVersion
     )
 
     applyDebugState(shim, resolvedNodule, esmResolver)
@@ -714,6 +715,10 @@ function tryGetVersion(shim) {
     return
   }
 
+  if (shim.pkgVersion) {
+    return shim.pkgVersion
+  }
+
   const packageInfo = shim.require('./package.json')
   if (!packageInfo) {
     return

test/unit/instrumentation/openai.test.js

@@ -19,7 +19,7 @@ test('openai unit tests', (t) => {
     agent.config.ai_monitoring = { enabled: true }
     const shim = new GenericShim(agent, 'openai')
     sandbox.stub(shim, 'require')
-    shim.require.returns({ version: '4.0.0' })
+    shim.pkgVersion = '4.0.0'
     sandbox.stub(shim.logger, 'debug')
 
     t.context.agent = agent
@@ -57,7 +57,7 @@ test('openai unit tests', (t) => {
   t.test('should not register instrumentation if openai is < 4.0.0', (t) => {
     const { shim, agent, initialize } = t.context
     const MockOpenAi = getMockModule()
-    shim.require.returns({ version: '3.7.0' })
+    shim.pkgVersion = '3.7.0'
     initialize(agent, MockOpenAi, 'openai', shim)
     t.equal(shim.logger.debug.callCount, 1, 'should log 2 debug messages')
     t.equal(

test/unit/llm-events/openai/chat-completion-summary.test.js

@@ -27,6 +27,7 @@ tap.test('LlmChatCompletionSummary', (t) => {
     helper.runInTransaction(agent, (tx) => {
       api.startSegment('fakeSegment', false, () => {
         const segment = api.shim.getActiveSegment()
+        segment.end()
         const chatSummaryEvent = new LlmChatCompletionSummary({
           agent,
           segment,

test/unit/llm-events/openai/common.js

@@ -58,7 +58,7 @@ function getExpectedResult(tx, event, type, completionId) {
     'ingest_source': 'Node'
   }
   const resKeys = {
-    'duration': trace.children[0].getExclusiveDurationInMillis(),
+    'duration': trace.children[0].getDurationInMillis(),
     'request.model': 'gpt-3.5-turbo-0613',
     'api_key_last_four_digits': 'sk-7890',
     'response.organization': 'new-relic',

test/unit/llm-events/openai/embedding.test.js

@@ -32,6 +32,7 @@ tap.test('LlmEmbedding', (t) => {
     helper.runInTransaction(agent, (tx) => {
       api.startSegment('fakeSegment', false, () => {
         const segment = api.shim.getActiveSegment()
+        segment.end()
         const embeddingEvent = new LlmEmbedding({ agent, segment, request: req, response: res })
         const expected = getExpectedResult(tx, embeddingEvent, 'embedding')
         t.same(embeddingEvent, expected)