wip: generateContentStreamInternal

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -5,7 +5,6 @@
 
 'use strict'
 
-const { geminiApiKey } = require('../../../lib/symbols')
 const {
   LlmChatCompletionMessage,
   LlmChatCompletionSummary,
@@ -17,6 +16,7 @@ const { extractLlmContext } = require('../../util/llm-utils')
 
 const { AI } = require('../../../lib/metrics/names')
 const { GEMINI } = AI
+const { DESTINATIONS } = require('../../config/attribute-filter')
 let TRACKING_METRIC = GEMINI.TRACKING_PREFIX
 
 /**
@@ -158,8 +158,6 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
   /**
    * 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) {
@@ -185,54 +183,76 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
     }
   )
 
-  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
-    }
+  /*
+    * 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
+    */
+
+  // TODO: might need to instrument processAfcStream too
+  // https://github.com/googleapis/js-genai/blob/cd0454862b4a0251d2606eeca8500b3b76004944/src/models.ts#L183
+  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
+      }
+      const [request] = args
 
-    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
+      return new RecorderSpec({
+        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 })
+              }
             }
-
-            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
           })
         }
-      }
+      })
     })
-  })
 
   /**
    * Instruments embedding creation
@@ -282,6 +302,4 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
       })
     }
   )
-
-  // TODO: shim.record generateContentStreamInternal
 }

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

@@ -15,8 +15,6 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
       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