generateContentStream and embedContent

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -7,10 +7,10 @@
 
 const { geminiApiKey } = require('../../../lib/symbols')
 const {
-    LlmChatCompletionMessage,
-    LlmChatCompletionSummary,
-    LlmEmbedding,
-    LlmErrorMessage
+  LlmChatCompletionMessage,
+  LlmChatCompletionSummary,
+  LlmEmbedding,
+  LlmErrorMessage
 } = require('../../../lib/llm-events/google-genai')
 const { RecorderSpec } = require('../../../lib/shim/specs')
 const { extractLlmContext } = require('../../util/llm-utils')
@@ -28,12 +28,12 @@ let TRACKING_METRIC = GEMINI.TRACKING_PREFIX
  * @param {object} params.msg LLM event
  */
 function recordEvent({ agent, type, msg }) {
-    const llmContext = extractLlmContext(agent)
+  const llmContext = extractLlmContext(agent)
 
-    agent.customEventAggregator.add([
-        { type, timestamp: Date.now() },
-        Object.assign({}, msg, llmContext)
-    ])
+  agent.customEventAggregator.add([
+    { type, timestamp: Date.now() },
+    Object.assign({}, msg, llmContext)
+  ])
 }
 
 /**
@@ -44,8 +44,8 @@ function recordEvent({ agent, type, msg }) {
  * @param {Transaction} params.transaction active transaction
  */
 function addLlmMeta({ agent, transaction }) {
-    agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount()
-    transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
+  agent.metrics.getOrCreateMetric(TRACKING_METRIC).incrementCallCount()
+  transaction.trace.attributes.addAttribute(DESTINATIONS.TRANS_EVENT, 'llm', true)
 }
 
 /**
@@ -65,103 +65,223 @@ function addLlmMeta({ agent, transaction }) {
  * @param {Transaction} params.transaction active transaction
  */
 function recordChatCompletionMessages({
+  agent,
+  shim,
+  segment,
+  request,
+  response,
+  err,
+  transaction
+}) {
+  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 = {}
+  }
+
+  // response.headers = segment[]
+  // explicitly end segment to consistent duration
+  // for both LLM events and the segment
+  segment.end()
+  const completionSummary = new LlmChatCompletionSummary({
     agent,
-    shim,
     segment,
+    transaction,
     request,
     response,
-    err,
-    transaction
-}) {
-    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 = {}
-    }
+    withError: err != null
+  })
 
-    // response.headers = segment[] 
-    // explicitly end segment to consistent duration
-    // for both LLM events and the segment
-    segment.end()
-    const completionSummary = new LlmChatCompletionSummary({
-        agent,
-        segment,
-        transaction,
-        request,
-        response,
-        withError: err != null
+  // Only take the first response message and append to input messages
+  const messages = [request.contents, response?.candidates?.[0]?.content?.parts?.[0]]
+  messages.forEach((message, index) => {
+    const completionMsg = new LlmChatCompletionMessage({
+      agent,
+      segment,
+      transaction,
+      request,
+      response,
+      index,
+      completionId: completionSummary.id,
+      message
     })
 
-    // Only take the first response message and append to input messages
-    const messages = [...request.messages, response?.choices?.[0]?.message]
-    messages.forEach((message, index) => {
-        const completionMsg = new LlmChatCompletionMessage({
+    recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: completionMsg })
+  })
+
+  recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: completionSummary })
+
+  if (err) {
+    const llmError = new LlmErrorMessage({ cause: err, summary: completionSummary, response })
+    agent.errors.add(transaction, err, llmError)
+  }
+
+  delete response.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.')
+    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
+    }
+  })
+
+  /**
+   * Instruments chat completion creation
+   * and creates the LLM events
+   *
+   * **Note**: Currently only for promises. streams will come later
+   */
+  shim.record(models.prototype, 'generateContentInternal',
+    function wrapGenerateContent(shim, func, name, args) {
+      const [request] = args
+
+      return new RecorderSpec({
+        name: GEMINI.COMPLETION,
+        promise: true,
+        after({ error: err, result: response, segment, transaction }) {
+          recordChatCompletionMessages({
             agent,
+            shim,
             segment,
             transaction,
             request,
             response,
-            index,
-            completionId: completionSummary.id,
-            message
-        })
+            err
+          })
+
+          addLlmMeta({ agent, transaction })
+        }
+      })
+    }
+  )
+
+  shim.record(models.prototype, 'generateContentStreamInternal', function wrapGenerateContentStream(shim, func, name, args) {
+    if (!agent.config.ai_monitoring.streaming.enabled) {
+      shim.logger.warn(
+        '`ai_monitoring.streaming.enabled` is set to `false`, stream will not be instrumented.'
+      )
+      agent.metrics.getOrCreateMetric(AI.STREAMING_DISABLED).incrementCallCount()
+      return
+    }
 
-        recordEvent({ agent, type: 'LlmChatCompletionMessage', msg: completionMsg })
+    shim.wrap(response, 'iterator', function wrapIterator(shim, orig) {
+      return async function * wrappedIterator() {
+        let content = ''
+        let role = ''
+        let chunk
+        let err
+        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 (streamErr) {
+          err = streamErr
+          throw err
+        } 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
+          })
+        }
+      }
     })
+  })
 
-    recordEvent({ agent, type: 'LlmChatCompletionSummary', msg: completionSummary })
+  /**
+   * Instruments embedding creation
+   * and creates LlmEmbedding event
+   */
+  shim.record(
+    models.prototype,
+    'embedContent',
+    function wrapEmbedContent(shim, func, name, args) {
+      const [request] = args
 
-    if (err) {
-        const llmError = new LlmErrorMessage({ cause: err, summary: completionSummary, response })
-        agent.errors.add(transaction, err, llmError)
-    }
+      return new RecorderSpec({
+        name: GEMINI.EMBEDDING,
+        promise: true,
+        after({ error: err, result: response, segment, transaction }) {
+          addLlmMeta({ agent, transaction })
 
-    delete response.headers
-}
+          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 = {}
+          }
 
-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.')
-        return
-    }
+          // explicitly end segment to get consistent duration
+          // for both LLM events and the segment
+          segment.end()
 
-    // 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
-    /**
-     * Instruments chat completion creation
-     * and creates the LLM events
-     *
-     * **Note**: Currently only for promises. streams will come later
-     */
-    shim.record(
-        models.prototype,
-        'generateContentInternal',
-        function wrapGenerateContent(shim, func, name, args) {
-            const [request] = args
-            const model = request?.model
-            const contents = request?.contents // the prompt
-
-            return new RecorderSpec({
-                name: GEMINI.COMPLETION,
-                promise: true,
-                after({ error: err, result: response, segment, transaction }) {
-                    recordChatCompletionMessages({
-                        agent,
-                        shim,
-                        segment,
-                        transaction,
-                        request,
-                        response,
-                        err
-                    })
-
-                    addLlmMeta({ agent, transaction })
-                }
-            })
+          const embedding = new LlmEmbedding({
+            agent,
+            segment,
+            transaction,
+            request,
+            response,
+            withError: err != null
+          })
+
+          recordEvent({ agent, type: 'LlmEmbedding', msg: embedding })
+
+          if (err) {
+            const llmError = new LlmErrorMessage({ cause: err, embedding, response })
+            shim.agent.errors.add(transaction, err, llmError)
+          }
+
+          // cleanup keys on response before returning to user code
+          // delete response.headers
         }
-    )
-}
+      })
+    }
+  )
+
+  // TODO: shim.record generateContentStreamInternal
+}

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

@@ -22,21 +22,16 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
     this.role = message?.role
     this.sequence = index
     this.completion_id = completionId
-    this.is_response = response?.choices?.[0]?.message?.content === message?.content
+    this.is_response = response?.candidates?.[0]?.content?.parts?.[0]?.text === message?.content
 
     if (agent.config.ai_monitoring.record_content.enabled === true) {
       this.content = message?.content
     }
 
-    const tokenCB = agent.llm?.tokenCountCallback
-    if (typeof tokenCB !== 'function') {
-      return
-    }
-
     if (this.is_response) {
-      this.token_count = tokenCB(this['response.model'], message?.content)
+      this.token_count = response.candidatesTokenCount
     } else {
-      this.token_count = tokenCB(request.model || request.engine, message?.content)
+      this.token_count = response.promptTokenCount
     }
   }
 }

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

@@ -10,9 +10,17 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
   constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
     super({ agent, segment, request, response, responseAttrs: true, transaction })
     this.error = withError
+    let requestMessagesLength = 0
+    if (Array.isArray(request?.contents)) {
+      requestMessagesLength = request.contents.length
+    } else if (typeof request?.contents === 'string') {
+      requestMessagesLength = 1
+    } else {
+      requestMessagesLength = 0
+    }
+    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['response.number_of_messages'] = request?.messages?.length + response?.choices?.length
-    this['response.choices.finish_reason'] = response?.choices?.[0]?.finish_reason
   }
 }