Address PR comments

Author: amychisholm03

lib/subscribers/google-genai/base.js

@@ -22,7 +22,7 @@ class GoogleGenAISubscriber extends Subscriber {
   }
 
   get enabled() {
-    return super.enabled && this.config?.ai_monitoring?.enabled
+    return super.enabled && this.agent?.config?.ai_monitoring?.enabled
   }
 
   /**
@@ -58,16 +58,11 @@ class GoogleGenAISubscriber extends Subscriber {
   recordChatCompletionMessages({
     segment,
     request,
-    response,
+    response = {},
     err,
     transaction
   }) {
     const agent = this.agent
-    if (!response) {
-      // If we get an error, it is possible that `response = null`.
-      // In that case, we define it to be an empty object.
-      response = {}
-    }
 
     // Explicitly end segment to provide consistent duration
     // for both LLM events and the segment
@@ -87,20 +82,21 @@ class GoogleGenAISubscriber extends Subscriber {
     const inputMessages = Array.isArray(request.contents) ? request.contents : [request.contents]
     const responseMessage = response?.candidates?.[0]?.content
     const messages = responseMessage !== undefined ? [...inputMessages, responseMessage] : inputMessages
-    messages.forEach((message, index) => {
+    for (let i = 0; i < messages.length; i++) {
+      const message = messages[i]
       const completionMsg = new LlmChatCompletionMessage({
         agent,
         segment,
         transaction,
         request,
         response,
-        index,
+        index: i,
         completionId: completionSummary.id,
         message
       })
 
       this.recordEvent({ type: 'LlmChatCompletionMessage', msg: completionMsg })
-    })
+    }
 
     this.recordEvent({ type: 'LlmChatCompletionSummary', msg: completionSummary })
 

lib/subscribers/google-genai/config.js

@@ -8,7 +8,7 @@ const generateContentInternal = {
   instrumentations: [
     {
       channelName: 'nr_generateContentInternal',
-      module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.cjs' },
+      module: { name: '@google/genai', versionRange: '>=1.1.0', filePath: 'dist/node/index.cjs' },
       functionQuery: {
         className: 'Models',
         methodName: 'generateContentInternal',
@@ -17,7 +17,7 @@ const generateContentInternal = {
     },
     {
       channelName: 'nr_generateContentInternal',
-      module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.mjs' },
+      module: { name: '@google/genai', versionRange: '>=1.1.0', filePath: 'dist/node/index.mjs' },
       functionQuery: {
         className: 'Models',
         methodName: 'generateContentInternal',
@@ -32,7 +32,7 @@ const generateContentStreamInternal = {
   instrumentations: [
     {
       channelName: 'nr_generateContentStreamInternal',
-      module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.cjs' },
+      module: { name: '@google/genai', versionRange: '>=1.1.0', filePath: 'dist/node/index.cjs' },
       functionQuery: {
         className: 'Models',
         methodName: 'generateContentStreamInternal',
@@ -41,7 +41,7 @@ const generateContentStreamInternal = {
     },
     {
       channelName: 'nr_generateContentStreamInternal',
-      module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.mjs' },
+      module: { name: '@google/genai', versionRange: '>=1.1.0', filePath: 'dist/node/index.mjs' },
       functionQuery: {
         className: 'Models',
         methodName: 'generateContentStreamInternal',
@@ -56,7 +56,7 @@ const embedContent = {
   instrumentations: [
     {
       channelName: 'nr_embedContent',
-      module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.cjs' },
+      module: { name: '@google/genai', versionRange: '>=1.1.0', filePath: 'dist/node/index.cjs' },
       functionQuery: {
         className: 'Models',
         methodName: 'embedContent',
@@ -65,7 +65,7 @@ const embedContent = {
     },
     {
       channelName: 'nr_embedContent',
-      module: { name: '@google/genai', versionRange: '>=1.1.0 <1.5.0 || >=1.5.1', filePath: 'dist/node/index.mjs' },
+      module: { name: '@google/genai', versionRange: '>=1.1.0', filePath: 'dist/node/index.mjs' },
       functionQuery: {
         className: 'Models',
         methodName: 'embedContent',

lib/subscribers/google-genai/embed-content.js

@@ -12,8 +12,8 @@ const {
 } = require('../../../lib/llm-events/google-genai')
 
 class GoogleGenAIEmbedContentSubscriber extends GoogleGenAISubscriber {
-  constructor ({ agent, logger, channelName = 'nr_embedContent' }) {
-    super({ agent, logger, channelName })
+  constructor ({ agent, logger }) {
+    super({ agent, logger, channelName: 'nr_embedContent' })
   }
 
   handler(data, ctx) {
@@ -32,15 +32,11 @@ class GoogleGenAIEmbedContentSubscriber extends GoogleGenAISubscriber {
       return
     }
     this.addLlmMeta({ transaction: ctx.transaction, version: data.moduleVersion })
-    let { result: response, arguments: args, error: err } = data
+    // If we get an error, it is possible that `response = null`.
+    // In that case, we define it to be an empty object.
+    const { result: response = {}, arguments: args, error: err } = data
     const [request] = args
 
-    if (!response) {
-      // If we get an error, it is possible that `response = null`.
-      // In that case, we define it to be an empty object.
-      response = {}
-    }
-
     // Explicitly end segment to get consistent duration
     // for both LLM events and the segment
     ctx.segment.end()

lib/subscribers/google-genai/generate-content-stream.js

@@ -7,8 +7,16 @@ const GoogleGenAIGenerateContentSubscriber = require('./generate-content')
 const { AI } = require('../../../lib/metrics/names')
 
 class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateContentSubscriber {
-  constructor({ agent, logger, channelName = 'nr_generateContentStreamInternal' }) {
-    super({ agent, logger, channelName })
+  constructor({ agent, logger }) {
+    super({ agent, logger, channelName: 'nr_generateContentStreamInternal' })
+  }
+
+  handler(data, ctx) {
+    if (!this.agent?.config?.ai_monitoring?.streaming?.enabled) {
+      return ctx
+    }
+
+    return super.handler(data, ctx)
   }
 
   /**
@@ -19,17 +27,7 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
    * @param {TraceSegment} params.segment the active trace segment
    * @param {Transaction} params.transaction the active transaction
    */
-  async instrumentStream({ request, response, segment, transaction }) {
-    const { agent, logger } = this
-
-    if (!agent.config.ai_monitoring.streaming.enabled) {
-      logger.warn(
-        '`ai_monitoring.streaming.enabled` is set to `false`, stream will not be instrumented.'
-      )
-      agent.metrics.getOrCreateMetric(AI.STREAMING_DISABLED).incrementCallCount()
-      return
-    }
-
+  instrumentStream({ request, response, segment, transaction }) {
     const originalNext = response.next
     let err
     let content
@@ -41,8 +39,8 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
       try {
         result = await originalNext.apply(response, nextArgs)
         if (result?.value?.text) {
-          modelVersion = result.value.modelVersion
-          content = result.value.candidates[0].content
+          modelVersion = result?.value?.modelVersion
+          content = result?.value?.candidates?.[0]?.content
           entireMessage += result.value.text // readonly variable that equates to result.value.candidates[0].content.parts[0].text
         }
         if (result?.value?.candidates?.[0]?.finishReason) {
@@ -88,19 +86,26 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
   }
 
   asyncEnd(data) {
+    if (!this.agent?.config?.ai_monitoring?.streaming?.enabled) {
+      this.logger.warn(
+        '`ai_monitoring.streaming.enabled` is set to `false`, stream will not be instrumented.'
+      )
+      this.agent.metrics.getOrCreateMetric(AI.STREAMING_DISABLED).incrementCallCount()
+
+      return
+    }
     const ctx = this.agent.tracer.getContext()
     if (!ctx?.segment || !ctx?.transaction) {
       return
     }
-    const { result: response, arguments: args, error: err } = data
+    const { result: response, arguments: args } = data
     const [request] = args
+
     this.instrumentStream({
       request,
-      headers: ctx.extras?.headers,
       response,
       segment: ctx.segment,
       transaction: ctx.transaction,
-      err
     })
 
     this.addLlmMeta({

lib/subscribers/google-genai/generate-content.js

@@ -13,6 +13,10 @@ class GoogleGenAIGenerateContentSubscriber extends GoogleGenAISubscriber {
   }
 
   handler(data, ctx) {
+    if (!this.agent?.config?.ai_monitoring?.enabled) {
+      return ctx
+    }
+
     const segment = this.agent.tracer.createSegment({
       name: GEMINI.COMPLETION,
       parent: ctx.segment,
@@ -22,6 +26,9 @@ class GoogleGenAIGenerateContentSubscriber extends GoogleGenAISubscriber {
   }
 
   asyncEnd(data) {
+    if (!this.agent?.config?.ai_monitoring?.enabled) {
+      return
+    }
     const ctx = this.agent.tracer.getContext()
     if (!ctx?.segment || !ctx?.transaction) {
       return
@@ -34,7 +41,6 @@ class GoogleGenAIGenerateContentSubscriber extends GoogleGenAISubscriber {
       transaction: ctx.transaction,
       request,
       response,
-      headers: ctx.extras?.headers,
       err
     })
 

test/versioned/google-genai/chat-completions.test.js

@@ -309,6 +309,30 @@ test('handles error in stream', (t, end) => {
 })
 
 // Other tests
+test('should not create llm events when ai_monitoring.enabled is false', (t, end) => {
+  const { client, agent } = t.nr
+  agent.config.ai_monitoring.enabled = false
+  helper.runInTransaction(agent, async (tx) => {
+    const model = 'gemini-2.0-flash'
+    const result = await client.models.generateContent({
+      model,
+      contents: 'You are a mathematician.'
+    })
+    assert.ok(result)
+
+    const events = agent.customEventAggregator.events.toArray()
+    assert.equal(events.length, 0, 'should not create llm events')
+
+    const activeSeg = agent.tracer.getSegment()
+    assert.equal(activeSeg?.isRoot, true)
+    const children = tx.trace.getChildren(activeSeg.id)
+    assert.notEqual(children?.[0]?.name, GEMINI.COMPLETION)
+
+    tx.end()
+    end()
+  })
+})
+
 test('should not create llm events when ai_monitoring.streaming.enabled is false', (t, end) => {
   const { client, agent } = t.nr
   agent.config.ai_monitoring.streaming.enabled = false
@@ -335,12 +359,17 @@ test('should not create llm events when ai_monitoring.streaming.enabled is false
     assert.equal(res, expectedRes.body.candidates[0].content.parts[0].text)
 
     const events = agent.customEventAggregator.events.toArray()
-    assert.equal(events.length, 0, 'should not llm events when streaming is disabled')
+    assert.equal(events.length, 0, 'should not create llm events when streaming is disabled')
     const streamingDisabled = agent.metrics.getOrCreateMetric(
       'Supportability/Nodejs/ML/Streaming/Disabled'
     )
     assert.equal(streamingDisabled.callCount > 0, true)
 
+    const activeSeg = agent.tracer.getSegment()
+    assert.equal(activeSeg?.isRoot, true)
+    const children = tx.trace.getChildren(activeSeg.id)
+    assert.notEqual(children?.[0]?.name, GEMINI.COMPLETION)
+
     tx.end()
     end()
   })