wip: generateContentStreamInternal

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -61,7 +61,7 @@ function addLlmMeta({ agent, transaction }) {
  * @param {Shim} params.shim the current shim instance
  * @param {TraceSegment} params.segment active segment from chat completion
  * @param {object} params.request chat completion params
- * @param {object} params.response chat completion response
+ * @param {GenerateContentResponse} params.response chat completion response
  * @param {boolean} [params.err] err if it exists
  * @param {Transaction} params.transaction active transaction
  */
@@ -127,6 +127,7 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
     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.
@@ -159,10 +160,6 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
   })
 
   const models = googleGenAi.Models
-  // TODO: why is generateContentInternal and generateContentStreamInternal
-  // exposed but not generateContent or generateContentStream?
-
-  // TODO: look at computeTokens and countTokens?
 
   /**
    * Instruments chat completion creation
@@ -192,14 +189,51 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
     }
   )
 
-  /*
-    * 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
-    */
+  // Wrap the stream to interept iterator.next calls
+  shim.wrap(models.prototype, 'generateContentStreamInternal', function wrapGenerateContentStream(shim, func) {
+    return async function wrappedGenerateContentStream(...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 iterator = await func.apply(this, args)
+
+      if (iterator && typeof iterator[Symbol.asyncIterator] === 'function') {
+        const originalNext = iterator.next
+        let err
+        let result
+
+        iterator.next = async function wrappedNext(...nextArgs) {
+          try {
+            result = await originalNext.apply(iterator, nextArgs)
+          } catch (streamErr) {
+            err = streamErr
+            throw err
+          } finally {
+          // TODO: How to integrate this with RecorderSpec?
+          // segment.touch()
+            recordChatCompletionMessages({
+              agent: shim.agent,
+              shim,
+              // segment,
+              // transaction,
+              // request,
+              response: result.value,
+            // err
+            })
+          }
+
+          return result
+        }
+      }
+
+      return iterator
+    }
+  })
 
   // TODO: might need to instrument processAfcStream too
   // https://github.com/googleapis/js-genai/blob/cd0454862b4a0251d2606eeca8500b3b76004944/src/models.ts#L183
@@ -212,23 +246,11 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
         agent.metrics.getOrCreateMetric(AI.STREAMING_DISABLED).incrementCallCount()
         return
       }
-      const [request] = args
 
       return new RecorderSpec({
         name: GEMINI.COMPLETION,
         promise: true,
-        after({ error: err, result: response, segment, transaction }) {
-          // TODO: actually need to handle the stream
-          recordChatCompletionMessages({
-            agent,
-            shim,
-            segment,
-            transaction,
-            request,
-            response,
-            err
-          })
-
+        after({ result: response, segment, transaction }) {
           addLlmMeta({ agent, transaction })
         }
       })
@@ -277,7 +299,7 @@ module.exports = function initialize(agent, googleGenAi, moduleName, shim) {
           }
 
           // cleanup keys on response before returning to user code
-          // delete response.headers
+          delete response.headers
         }
       })
     }