wip: tests

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -5,6 +5,7 @@
 
 'use strict'
 
+const { geminiApiKey, googleGenAiHeaders } = require('../../../lib/symbols')
 const {
   LlmChatCompletionMessage,
   LlmChatCompletionSummary,
@@ -119,42 +120,67 @@ function recordChatCompletionMessages({
   delete response.headers
 }
 
+/**
+ * Adds apiKey and response headers to the active segment
+ * on symbols
+ *
+ * @param {object} params input params
+ * @param {Shim} params.shim instance of shim
+ * @param {object} params.result from openai request
+ * @param {string} params.apiKey api key from openai client
+ */
+function decorateSegment({ shim, result, apiKey }) {
+  const segment = shim.getActiveSegment()
+
+  if (segment) {
+    segment[geminiApiKey] = apiKey
+
+    const headers = result?.response?.headers
+      ? Object.fromEntries(result.response.headers)
+      : { ...result?.headers }
+    segment[googleGenAiHeaders] = headers
+  }
+}
+
 module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
   if (agent?.config?.ai_monitoring?.enabled !== true) {
-    shim.logger.debug('config.ai_monitoring.enabled is set to false. Skipping instrumentation.')
+    shim.logger.debug('config.ai_monitoring.enabled is set to false.')
     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}`
 
-  const models = googleGenAi.Models
-  // TODO: why is generateContentInternal and generateContentStreamInternal
-  // exposed but not generateContent or generateContentStream?
-
   /**
    * Instrumentation is only done to get the response headers and attach
    * to the active segment as @google/genai hides the headers from the functions
    * we are trying to instrument.
    * see: https://github.com/googleapis/js-genai/blob/cd0454862b4a0251d2606eeca8500b3b76004944/src/models.ts#L200
-   *
-   * TODO: Do we even need the headers?
    */
-  shim.wrap(models.prototype, 'processParamsForMcpUsage', function wrapProcessParamsForMcpUsage(shim, original) {
-    return async function wrappedProcessParamsForMcpUsage(...args) {
-      // Call the original function and capture the result
-      const newParams = await original.apply(this, arguments)
-
-      // Inspect the headers in newParams
-      const headers = newParams?.config?.httpOptions?.headers
-      shim.logger.debug('Headers in newParams:', headers)
-
-      // Return the modified newParams
-      return newParams
+  const httpResponse = googleGenAi.HttpResponse
+  shim.wrap(httpResponse.prototype, 'json', function wrapJson(shim, func) {
+    return async function wrappedJson() {
+      const response = func.apply(this, arguments)
+      if (await response) {
+        // TODO: this does get some headers but not 'x-goog*'
+        const headers = this.headers
+        if (headers) {
+          decorateSegment({
+            shim,
+            result: response,
+            apiKey: this[geminiApiKey]
+          })
+        }
+      }
+      return response
     }
   })
 
+  const models = googleGenAi.Models
+  // TODO: why is generateContentInternal and generateContentStreamInternal
+  // exposed but not generateContent or generateContentStream?
+
   /**
    * Instruments chat completion creation
    * and creates the LLM events
@@ -209,47 +235,18 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
         name: GEMINI.COMPLETION,
         promise: true,
         after({ error: err, result: response, segment, transaction }) {
-          // Symbol.asyncIterator
-          // FIXME: it's causing recursion
-          shim.wrap(response, Symbol.asyncIterator, function wrapIterator(shim, orig) {
-            const originalAsyncIterator = orig
-            return async function * wrappedIterator() {
-              let content = ''
-              let role = ''
-              let chunk
-              let err
-              try {
-                const iterator = originalAsyncIterator.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 (streamErr) {
-                err = streamErr
-              } 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({
-                  agent: shim.agent,
-                  shim,
-                  segment,
-                  transaction,
-                  request,
-                  response: chunk,
-                  err
-                })
-
-                addLlmMeta({ agent, transaction })
-              }
-            }
+          // TODO: actually need to handle the stream
+          recordChatCompletionMessages({
+            agent,
+            shim,
+            segment,
+            transaction,
+            request,
+            response,
+            err
           })
+
+          addLlmMeta({ agent, transaction })
         }
       })
     })

lib/llm-events/google-genai/chat-completion-message.js

@@ -18,7 +18,7 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
     transaction
   }) {
     super({ agent, segment, request, response, transaction })
-    this.id = `${response.id}-${index}`
+    this.id = `${response.responseId}-${index}`
     this.role = message?.role
     this.sequence = index
     this.completion_id = completionId

lib/llm-events/google-genai/chat-completion-summary.js

@@ -18,7 +18,7 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
     }
     this['response.number_of_messages'] = requestMessagesLength + (response?.candidates?.length || 0)
     this['response.choices.finish_reason'] = response?.candidates?.[0]?.finishReason
-    this['request.max_tokens'] = request.max_tokens
-    this['request.temperature'] = request.temperature
+    this['request.max_tokens'] = request.config?.maxOutputTokens
+    this['request.temperature'] = request.config?.temperature
   }
 }

lib/llm-events/google-genai/embedding.js

@@ -20,19 +20,18 @@ const defaultParams = {}
 class LlmEmbedding extends LlmEvent {
   constructor(params = defaultParams) {
     super(params)
-    const { agent, input } = params
-    const tokenCb = agent?.llm?.tokenCountCallback
+    const { agent, request } = params
 
+    // TODO: idk if this is correct for input
     this.input = agent.config?.ai_monitoring?.record_content?.enabled
-      ? input
+      ? request?.contents
       : undefined
     this.error = params.isError
     this.duration = params.segment.getDurationInMillis()
-
-    // Even if not recording content, we should use the local token counting callback to record token usage
-    if (typeof tokenCb === 'function') {
-      this.token_count = tokenCb(this.bedrockCommand.modelId, input)
-    }
+    // TODO: idk if this is correct for token count
+    this.token_count = Array.isArray(params.response?.embeddings)
+      ? params.response.embeddings.reduce((sum, e) => sum + (e?.values?.length || 0), 0)
+      : undefined
   }
 }
 

lib/llm-events/google-genai/event.js

@@ -14,11 +14,11 @@ module.exports = class LlmEvent extends BaseEvent {
 
     this.id = makeId(36)
     this.appName = agent.config.applications()[0]
-    this.request_id = response?.headers?.['x-request-id']
+    this.request_id = response?.headers?.['x-request-id'] // TODO: what is this in Gemini?
     this.trace_id = transaction?.traceId
     this.span_id = segment?.id
-    this['response.model'] = response.modelVersion
-    this['request.model'] = request.model
+    this['response.model'] = response?.modelVersion
+    this['request.model'] = request?.model
     this.vendor = 'gemini'
     this.ingest_source = 'Node'
     this.metadata = agent
@@ -37,6 +37,7 @@ module.exports = class LlmEvent extends BaseEvent {
 
   responseAttrs(response) {
     // TODO: no response.headers?
+    // relevant headers will be prefixed with 'x-goog'?
     this['response.organization'] = response?.headers?.['google-organization']
     this['response.headers.llmVersion'] = response?.headers?.['gemini-version']
     this['response.headers.ratelimitLimitRequests'] =

lib/symbols.js

@@ -25,6 +25,7 @@ module.exports = {
   openAiHeaders: Symbol('openAiHeaders'),
   openAiApiKey: Symbol('openAiApiKey'),
   geminiApiKey: Symbol('geminiApiKey'),
+  googleGenAiHeaders: Symbol('googleGenAiHeaders'),
   parentSegment: Symbol('parentSegment'),
   langchainRunId: Symbol('runId'),
   otelSynthesis: Symbol('otelSynthesis'),

test/unit/instrumentation/@google/genai.test.js

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const test = require('node:test')
+const assert = require('node:assert')
+const helper = require('../../../lib/agent_helper')
+const GenericShim = require('../../../../lib/shim/shim')
+const sinon = require('sinon')
+
+test('@google/genai unit.tests', async (t) => {
+  t.beforeEach(function (ctx) {
+    ctx.nr = {}
+    const sandbox = sinon.createSandbox()
+    const agent = helper.loadMockedAgent()
+    agent.config.ai_monitoring = { enabled: true, streaming: { enabled: true } }
+    const shim = new GenericShim(agent, '@google/genai')
+    sandbox.stub(shim.logger, 'debug')
+    sandbox.stub(shim.logger, 'warn')
+
+    ctx.nr.agent = agent
+    ctx.nr.shim = shim
+    ctx.nr.sandbox = sandbox
+    ctx.nr.initialize = require('../../../../lib/instrumentation/@google/genai.js')
+  })
+
+  t.afterEach(function (ctx) {
+    helper.unloadAgent(ctx.nr.agent)
+    ctx.nr.sandbox.restore()
+  })
+
+  function getMockModule() {
+    function GoogleGenAi() {}
+    GoogleGenAi.HttpResponse = function () {}
+    GoogleGenAi.HttpResponse.prototype.json = async function () {}
+    GoogleGenAi.Models = function () {}
+    GoogleGenAi.Models.prototype.generateContentInternal = async function () {}
+    GoogleGenAi.Models.prototype.generateContentStreamInternal = async function () {}
+    GoogleGenAi.Models.prototype.embedContent = async function () {}
+    return GoogleGenAi
+  }
+
+  await t.test('should instrument @google/genai', (t, end) => {
+    const { shim, agent, initialize } = t.nr
+    const MockGoogleGenAi = getMockModule()
+    initialize(agent, MockGoogleGenAi, '@google/genai', shim)
+    assert.equal(shim.logger.debug.callCount, 0, 'should not log debug messages')
+    const isWrapped = shim.isWrapped(MockGoogleGenAi.Models.prototype.generateContentInternal)
+    assert.equal(isWrapped, true, 'should wrap models generateContentInternal')
+    end()
+  })
+
+  await t.test(
+    'should not instrument generate content streams if ai_monitoring.streaming.enabled is false',
+    (t, end) => {
+      const { shim, agent, initialize } = t.nr
+      agent.config.ai_monitoring.streaming.enabled = false
+      const MockGoogleGenAi = getMockModule()
+      initialize(agent, MockGoogleGenAi, '@google/genai', shim)
+      const models = new MockGoogleGenAi.Models()
+
+      helper.runInTransaction(agent, async () => {
+        await models.generateContentStreamInternal()
+        assert.equal(
+          shim.logger.warn.args[0][0],
+          '`ai_monitoring.streaming.enabled` is set to `false`, stream will not be instrumented.'
+        )
+        end()
+      })
+    }
+  )
+
+  await t.test(
+    'should not register instrumentation if ai_monitoring.enabled is false',
+    (t, end) => {
+      const { shim, agent, initialize } = t.nr
+      const MockGoogleGenAi = getMockModule()
+      agent.config.ai_monitoring = { enabled: false }
+
+      initialize(agent, MockGoogleGenAi, '@google/genai', shim)
+      assert.equal(shim.logger.debug.callCount, 1, 'should log 1 debug message')
+      assert.equal(shim.logger.debug.args[0][0], 'config.ai_monitoring.enabled is set to false.')
+      const isWrapped = shim.isWrapped(MockGoogleGenAi.Models.prototype.generateContentInternal)
+      assert.equal(isWrapped, false, 'should not wrap models generate content internal')
+      end()
+    }
+  )
+})

test/unit/llm-events/google-genai/chat-completion-message.test.js

@@ -3,4 +3,42 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-// TODO
+'use strict'
+
+const test = require('node:test')
+const assert = require('node:assert')
+const LlmChatCompletionMessage = require('../../../../lib/llm-events/google-genai/chat-completion-message')
+const helper = require('../../../lib/agent_helper')
+const { req, res, getExpectedResult } = require('./common')
+
+test.beforeEach((ctx) => {
+  ctx.nr = {}
+  ctx.nr.agent = helper.loadMockedAgent()
+})
+
+test.afterEach((ctx) => {
+  helper.unloadAgent(ctx.nr.agent)
+})
+
+test('should create a LlmChatCompletionMessage event', (t, end) => {
+  const { agent } = t.nr
+  const api = helper.getAgentApi()
+  helper.runInTransaction(agent, (tx) => {
+    api.startSegment('fakeSegment', false, () => {
+      const segment = api.shim.getActiveSegment()
+      const chatMessageEvent = new LlmChatCompletionMessage({
+        transaction: tx,
+        agent,
+        segment,
+        request: req,
+        response: res,
+        message: req.contents,
+        index: 0
+      })
+      // TODO: add the expected result
+      const expected = getExpectedResult(tx, {}, 'message')
+      assert.deepEqual(chatMessageEvent, expected)
+      end()
+    })
+  })
+})