fix: handle case when content is undefined in candidate.

PiperOrigin-RevId: 651223704
googleapis · Jul 11, 2024 · f16f040 · f16f040
1 parent 23c9b68
commit f16f040
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 1 deletion.
diff --git a/src/functions/post_fetch_processing.ts b/src/functions/post_fetch_processing.ts
@@ -204,7 +204,7 @@ export function aggregateResponses(
         aggregatedResponse.candidates[i] = {
           index: response.candidates[i].index ?? i,
           content: {
-            role: response.candidates[i].content.role ?? constants.MODEL_ROLE,
+            role: response.candidates[i].content?.role ?? constants.MODEL_ROLE,
             parts: [{text: ''}],
           },
         } as GenerateContentCandidate;

diff --git a/src/functions/test/post_fetch_processing_test.ts b/src/functions/test/post_fetch_processing_test.ts
@@ -19,10 +19,12 @@ import {
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_1,
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_2,
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_3,
+  AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_4,
   COUNT_TOKENS_RESPONSE_1,
   STREAM_RESPONSE_CHUNKS_1,
   STREAM_RESPONSE_CHUNKS_2,
   STREAM_RESPONSE_CHUNKS_3,
+  STREAM_RESPONSE_CHUNKS_4,
   UNARY_RESPONSE_1,
   UNARY_RESPONSE_MISSING_ROLE_INDEX,
 } from './test_data';
@@ -74,6 +76,14 @@ describe('aggregateResponses', () => {
       JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_3)
     );
   });
+
+  it('missing content, should add role and return empty content', () => {
+    const actualResult = aggregateResponses(STREAM_RESPONSE_CHUNKS_4);
+
+    expect(JSON.stringify(actualResult)).toEqual(
+      JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_4)
+    );
+  });
 });
 
 describe('processUnary', () => {

diff --git a/src/functions/test/test_data.ts b/src/functions/test/test_data.ts
@@ -969,6 +969,26 @@ export const AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_3: GenerateContentRespon
       },
     ],
   } as GenerateContentResponse;
+
+export const STREAM_RESPONSE_CHUNKS_4: GenerateContentResponse[] = [
+  {
+    candidates: [{}],
+  },
+] as GenerateContentResponse[];
+
+export const AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_4: GenerateContentResponse =
+  {
+    candidates: [
+      {
+        index: 0,
+        content: {
+          role: 'model',
+          parts: [{text: ''}],
+        },
+      },
+    ],
+  } as GenerateContentResponse;
+
 export const UNARY_RESPONSE_1: GenerateContentResponse = {
   candidates: [
     {