fix: aggregate text & functionCall(s) correctly (#497)

* fix: aggregate text & functionCall(s) correctly

content parts should have either text or functionCall set, but never both. When aggregating, combine the text (as before), but add functionCall(s) as new parts. Then when done, if the text from part[0] is empty, that part can be removed.

fixes #494

* fix: aggregate text & functionCall(s) correctly

don't delete the text part, because it could just be added back again when working with multiple responses (and the end up as `"undefined\n"` for example). Leave the text part, and always push functions as new parts.

fixes #494

* fix: aggregate text & functionCall(s) correctly

the AI API gives an error for empty text, so the empty text has to be deleted after aggregating all responses

fixes #494

* chore: fix style using script

---------

Co-authored-by: Yvonne Yu <150068659+yyyu-google@users.noreply.github.com>

Author: 6utt3rfly

src/functions/post_fetch_processing.ts

@@ -249,16 +249,15 @@ export function aggregateResponses(
         response.candidates[i].content.parts &&
         response.candidates[i].content.parts.length > 0
       ) {
+        const {parts} = aggregatedResponse.candidates[i].content;
         for (const part of response.candidates[i].content.parts) {
+          // NOTE: cannot have text and functionCall both in the same part.
+          // add functionCall(s) to new parts.
           if (part.text) {
-            aggregatedResponse.candidates[i].content.parts[0].text += part.text;
+            parts[0].text += part.text;
           }
           if (part.functionCall) {
-            aggregatedResponse.candidates[i].content.parts[0].functionCall =
-              part.functionCall;
-            // the empty 'text' key should be removed if functionCall is in the
-            // response
-            delete aggregatedResponse.candidates[i].content.parts[0].text;
+            parts.push({functionCall: part.functionCall});
           }
         }
       }
@@ -273,6 +272,16 @@ export function aggregateResponses(
       }
     }
   }
+  if (aggregatedResponse.candidates?.length) {
+    aggregatedResponse.candidates.forEach(candidate => {
+      if (
+        candidate.content.parts.length > 1 &&
+        candidate.content.parts[0].text === ''
+      ) {
+        candidate.content.parts.shift(); // remove empty text parameter
+      }
+    });
+  }
   return aggregatedResponse;
 }
 

src/functions/test/post_fetch_processing_test.ts

@@ -20,11 +20,15 @@ import {
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_2,
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_3,
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_4,
+  AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_5,
+  AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_6,
   COUNT_TOKENS_RESPONSE_1,
   STREAM_RESPONSE_CHUNKS_1,
   STREAM_RESPONSE_CHUNKS_2,
   STREAM_RESPONSE_CHUNKS_3,
   STREAM_RESPONSE_CHUNKS_4,
+  STREAM_RESPONSE_CHUNKS_5,
+  STREAM_RESPONSE_CHUNKS_6,
   UNARY_RESPONSE_1,
   UNARY_RESPONSE_MISSING_ROLE_INDEX,
 } from './test_data';
@@ -84,6 +88,22 @@ describe('aggregateResponses', () => {
       JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_4)
     );
   });
+
+  it('should aggregate text and functionCalls to separate parts', () => {
+    const actualResult = aggregateResponses(STREAM_RESPONSE_CHUNKS_5);
+
+    expect(JSON.stringify(actualResult)).toEqual(
+      JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_5)
+    );
+  });
+
+  it('should aggregate functionCalls with the empty text removed', () => {
+    const actualResult = aggregateResponses(STREAM_RESPONSE_CHUNKS_6);
+
+    expect(JSON.stringify(actualResult)).toEqual(
+      JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_6)
+    );
+  });
 });
 
 describe('processUnary', () => {

src/functions/test/test_data.ts

@@ -905,6 +905,71 @@ export const AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_4: GenerateContentRespon
     ],
   } as GenerateContentResponse;
 
+export const STREAM_RESPONSE_CHUNKS_5: GenerateContentResponse[] = [
+  {
+    candidates: [
+      {
+        content: {
+          parts: [
+            {functionCall: {name: 'fn', args: {a: 1}}},
+            {text: 'chunk1Text1'},
+          ],
+        },
+      },
+    ],
+  },
+  {
+    candidates: [
+      {
+        content: {
+          parts: [{text: 'chunk2Text1'}],
+        },
+      },
+    ],
+  },
+] as GenerateContentResponse[];
+
+export const AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_5: GenerateContentResponse =
+  {
+    candidates: [
+      {
+        index: 0,
+        content: {
+          role: 'model',
+          parts: [
+            {text: 'chunk1Text1chunk2Text1'},
+            {functionCall: {name: 'fn', args: {a: 1}}},
+          ],
+        },
+      },
+    ],
+  } as GenerateContentResponse;
+
+export const STREAM_RESPONSE_CHUNKS_6: GenerateContentResponse[] = [
+  {
+    candidates: [
+      {
+        content: {
+          parts: [{functionCall: {name: 'fn', args: {a: 1}}}],
+        },
+      },
+    ],
+  },
+] as GenerateContentResponse[];
+
+export const AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_6: GenerateContentResponse =
+  {
+    candidates: [
+      {
+        index: 0,
+        content: {
+          role: 'model',
+          parts: [{functionCall: {name: 'fn', args: {a: 1}}}],
+        },
+      },
+    ],
+  } as GenerateContentResponse;
+
 export const UNARY_RESPONSE_1: GenerateContentResponse = {
   candidates: [
     {