fix: correct GenerateContentCandidate interface and GenerateContentResponse interface

PiperOrigin-RevId: 617875952

Author: yyyu-google

src/functions/post_fetch_processing.ts

@@ -16,8 +16,8 @@
  */
 
 import {
-  Citation,
   CitationMetadata,
+  Content,
   CountTokensResponse,
   GenerateContentCandidate,
   GenerateContentResponse,
@@ -26,6 +26,7 @@ import {
   Part,
   StreamGenerateContentResult,
 } from '../types/content';
+import {constants} from '../util';
 import {ClientError, GoogleGenerativeAIError} from '../types/errors';
 
 export async function throwErrorIfNotOK(response: Response | undefined) {
@@ -57,7 +58,7 @@ async function* generateResponseSequence(
     if (done) {
       break;
     }
-    yield addCandidateFunctionCalls(value);
+    yield addMissingFields(value);
   }
 }
 
@@ -183,18 +184,28 @@ export function aggregateResponses(
     );
   }
 
-  const aggregatedResponse: GenerateContentResponse = {
-    candidates: [],
-    promptFeedback: lastResponse.promptFeedback,
-    usageMetadata: lastResponse.usageMetadata,
-  };
+  const aggregatedResponse: GenerateContentResponse = {};
+
+  if (lastResponse.promptFeedback) {
+    aggregatedResponse.promptFeedback = lastResponse.promptFeedback;
+  }
+  if (lastResponse.usageMetadata) {
+    aggregatedResponse.usageMetadata = lastResponse.usageMetadata;
+  }
+
   for (const response of responses) {
+    if (!response.candidates || response.candidates.length === 0) {
+      continue;
+    }
     for (let i = 0; i < response.candidates.length; i++) {
+      if (!aggregatedResponse.candidates) {
+        aggregatedResponse.candidates = [];
+      }
       if (!aggregatedResponse.candidates[i]) {
         aggregatedResponse.candidates[i] = {
-          index: response.candidates[i].index,
+          index: response.candidates[i].index ?? i,
           content: {
-            role: response.candidates[i].content.role,
+            role: response.candidates[i].content.role ?? constants.MODEL_ROLE,
             parts: [{text: ''}],
           },
         } as GenerateContentCandidate;
@@ -246,8 +257,6 @@ export function aggregateResponses(
       }
     }
   }
-  aggregatedResponse.promptFeedback =
-    responses[responses.length - 1].promptFeedback;
   return aggregatedResponse;
 }
 
@@ -304,6 +313,33 @@ function aggregateGroundingMetadataForCandidate(
   return groundingMetadataAggregated;
 }
 
+function addMissingIndexAndRole(
+  response: GenerateContentResponse
+): GenerateContentResponse {
+  const generateContentResponse = response as GenerateContentResponse;
+  if (
+    generateContentResponse.candidates &&
+    generateContentResponse.candidates.length > 0
+  ) {
+    generateContentResponse.candidates.forEach((candidate, index) => {
+      if (candidate.index === undefined) {
+        generateContentResponse.candidates![index].index = index;
+      }
+
+      if (candidate.content === undefined) {
+        generateContentResponse.candidates![index].content = {} as Content;
+      }
+
+      if (candidate.content.role === undefined) {
+        generateContentResponse.candidates![index].content.role =
+          constants.MODEL_ROLE;
+      }
+    });
+  }
+
+  return generateContentResponse;
+}
+
 function addCandidateFunctionCalls(
   response: GenerateContentResponse
 ): GenerateContentResponse {
@@ -328,6 +364,13 @@ function addCandidateFunctionCalls(
   return response;
 }
 
+function addMissingFields(
+  response: GenerateContentResponse
+): GenerateContentResponse {
+  const generateContentResponse = addMissingIndexAndRole(response);
+  return addCandidateFunctionCalls(generateContentResponse);
+}
+
 /**
  * Process model responses from generateContent
  * @ignore
@@ -338,8 +381,9 @@ export async function processUnary(
   if (response !== undefined) {
     // ts-ignore
     const responseJson = await response.json();
+    const generateContentResponse = addMissingIndexAndRole(responseJson);
     return Promise.resolve({
-      response: addCandidateFunctionCalls(responseJson),
+      response: addCandidateFunctionCalls(generateContentResponse),
     });
   }
 

src/functions/test/functions_test.ts

@@ -585,7 +585,7 @@ describe('generateContent', () => {
       TEST_API_ENDPOINT
     );
     expect(
-      resp.response.candidates[0].citationMetadata?.citations.length
+      resp.response.candidates![0].citationMetadata?.citations.length
     ).toEqual(
       TEST_MODEL_RESPONSE.candidates[0].citationMetadata.citations.length
     );
@@ -617,9 +617,9 @@ describe('generateContent', () => {
       TEST_API_ENDPOINT
     );
     expect(actualResult).toEqual(expectedResult);
-    expect(actualResult.response.candidates[0].functionCalls).toHaveSize(1);
-    expect(actualResult.response.candidates[0].functionCalls).toEqual([
-      expectedResult.response.candidates[0].content.parts[0].functionCall!,
+    expect(actualResult.response.candidates![0].functionCalls).toHaveSize(1);
+    expect(actualResult.response.candidates![0].functionCalls).toEqual([
+      expectedResult.response.candidates![0].content.parts[0].functionCall!,
     ]);
   });
 
@@ -652,7 +652,9 @@ describe('generateContent', () => {
       TEST_API_ENDPOINT
     );
     expect(actualResult).toEqual(expectedResult);
-    expect(actualResult.response.candidates[0].functionCalls).not.toBeDefined();
+    expect(
+      actualResult.response.candidates![0].functionCalls
+    ).not.toBeDefined();
   });
 
   it('returns empty candidates when response is empty', async () => {

src/functions/test/post_fetch_processing_test.ts

@@ -18,10 +18,13 @@
 import {
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_1,
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_2,
+  AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_3,
   COUNT_TOKENS_RESPONSE_1,
   STREAM_RESPONSE_CHUNKS_1,
   STREAM_RESPONSE_CHUNKS_2,
+  STREAM_RESPONSE_CHUNKS_3,
   UNARY_RESPONSE_1,
+  UNARY_RESPONSE_MISSING_ROLE_INDEX,
 } from './test_data';
 import * as PostFetchFunctions from '../post_fetch_processing';
 import {aggregateResponses} from '../post_fetch_processing';
@@ -60,6 +63,18 @@ describe('aggregateResponses', () => {
       JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_2)
     );
   });
+
+  it('missing candidates, should return {}', () => {
+    expect(aggregateResponses([{}, {}])).toEqual({});
+  });
+
+  it('missing role and index, should add role and index', () => {
+    const actualResult = aggregateResponses(STREAM_RESPONSE_CHUNKS_3);
+
+    expect(JSON.stringify(actualResult)).toEqual(
+      JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_3)
+    );
+  });
 });
 
 describe('processUnary', () => {
@@ -80,6 +95,41 @@ describe('processUnary', () => {
 
     expect(actualResponse).toEqual(UNARY_RESPONSE_1);
   });
+
+  it('response missing role and index, should add role and index', async () => {
+    const fetchResult = new Response(
+      JSON.stringify(UNARY_RESPONSE_MISSING_ROLE_INDEX),
+      fetchResponseObj
+    );
+    const expectedResult = UNARY_RESPONSE_1;
+    spyOn(global, 'fetch').and.resolveTo(fetchResult);
+    const actualResult = await generateContent(
+      LOCATION,
+      PROJECT,
+      PUBLISHER_MODEL_ENDPOINT,
+      TOKEN,
+      GENERATE_CONTENT_REQUEST
+    );
+    const actualResponse = actualResult.response;
+
+    expect(actualResponse).toEqual(expectedResult);
+  });
+
+  it('candidate undefined, should return empty response', async () => {
+    spyOn(global, 'fetch').and.resolveTo(
+      new Response(JSON.stringify({}), fetchResponseObj)
+    );
+    const actualResult = await generateContent(
+      LOCATION,
+      PROJECT,
+      PUBLISHER_MODEL_ENDPOINT,
+      TOKEN,
+      GENERATE_CONTENT_REQUEST
+    );
+    const actualResponse = actualResult.response;
+
+    expect(actualResponse).toEqual({});
+  });
 });
 
 describe('processStream', () => {