concat stream response together

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -94,6 +94,8 @@ function recordChatCompletionMessages({
   })
 
   // Only take the first response message and append to input messages
+  // request.contents can be a string or an array of strings
+  // reponse.candidates is an array of candidates (choices), we only take the first one
   const inputMessages = Array.isArray(request.contents) ? request.contents : [request.contents]
   const responseMessage = response?.candidates?.[0]?.content
   const messages = responseMessage !== undefined ? [...inputMessages, responseMessage] : inputMessages
@@ -130,29 +132,56 @@ function instrumentStream ({ agent, shim, request, response, segment, transactio
     agent.metrics.getOrCreateMetric(AI.STREAMING_DISABLED).incrementCallCount()
   }
 
-  let err
+  let err = false
+  let content
+  let modelVersion
+  let finishReason
+  let entireMessage = ''
   shim.wrap(response, 'next', function wrapNext(shim, originalNext) {
     return async function wrappedNext(...nextArgs) {
-      let result
+      let result = {}
       try {
         result = await originalNext.apply(response, nextArgs)
+        if (result?.value?.candidates?.[0]?.content) {
+          modelVersion = result.value.modelVersion
+          content = result.value.candidates[0].content
+          entireMessage += content.parts[0].text
+        }
+        if (result?.value?.candidates?.[0]?.finishReason) {
+          finishReason = result.value.candidates[0].finishReason
+        }
       } catch (streamErr) {
         err = streamErr
         throw err
       } finally {
-        // update segment duration since we want to extend the time it took to
-        // handle the stream
+        // 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: result?.value,
-          err
-        })
+        if (result?.done || err) {
+          // result will be {value: undefined, done: true}
+          // when the stream is done, so we need to create
+          // a final response with the entire content
+          //
+          // also need to enter this block if there was an
+          // error, so we can record it
+          content.parts[0].text = entireMessage
+          result.value = {
+            candidates: [
+              { content, finishReason }
+            ],
+            modelVersion
+          }
+          recordChatCompletionMessages({
+            agent: shim.agent,
+            shim,
+            segment,
+            transaction,
+            request,
+            response: result.value,
+            err
+          })
+        }
       }
       return result
     }

lib/symbols.js

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