add streaming bug fix

Author: amychisholm03

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

@@ -29,23 +29,26 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
    */
   instrumentStream({ request, response, segment, transaction }) {
     const originalNext = response.next
+    let isDone = false
+    let cachedResult = {}
     let err
-    let content
-    let modelVersion
-    let finishReason
     let entireMessage = ''
     response.next = async function wrappedNext(...nextArgs) {
       let result = {}
       try {
         result = await originalNext.apply(response, nextArgs)
+        // When the stream is done we get {value: undefined, done: true}
+        // we need to cache the composed value and add the entire message
+        // back in later
+        if (result.done === true) {
+          isDone = true
+        } else {
+          cachedResult = result.value
+        }
+
         if (result?.value?.text) {
-          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) {
-          finishReason = result.value.candidates[0].finishReason
-        }
       } catch (streamErr) {
         err = streamErr
         throw err
@@ -54,29 +57,17 @@ class GoogleGenAIGenerateContentStreamSubscriber extends GoogleGenAIGenerateCont
         // time it took to handle the stream
         segment.touch()
 
-        // result will be {value: undefined, done: true}
-        // when the stream is done, so we need to create
-        // a mock GenerateContentResponse object with
-        // the entire message
-        //
         // also need to enter this block if there was an
         // error, so we can record it
-        if (result?.done || err) {
-          if (content) {
-            content.parts[0].text = entireMessage
-            result.value = {
-              candidates: [
-                { content, finishReason }
-              ],
-              modelVersion
-            }
+        if (isDone || err) {
+          if (cachedResult?.candidates?.[0]?.content?.parts) {
+            cachedResult.candidates[0].content.parts[0].text = entireMessage
           }
-
           this.recordChatCompletionMessages({
             segment,
             transaction,
             request,
-            response: result?.value,
+            response: cachedResult,
             err
           })
         }