fix: fix bug in safetyRatings handling, fix incomplete content interfaces, and add unit test for stream response handling

PiperOrigin-RevId: 616867948

Author: yyyu-google

src/functions/post_fetch_processing.ts

@@ -117,7 +117,7 @@ async function getResponsePromise(
  * GenerateContentResponse in each iteration.
  * @ignore
  */
-export function getResponseStream(
+function getResponseStream(
   inputStream: ReadableStream<string>
 ): ReadableStream<unknown> {
   const reader = inputStream.getReader();
@@ -170,8 +170,9 @@ export function getResponseStream(
  * Aggregates an array of `GenerateContentResponse`s into a single
  * GenerateContentResponse.
  * @ignore
+ * @VisibleForTesting
  */
-function aggregateResponses(
+export function aggregateResponses(
   responses: GenerateContentResponse[]
 ): GenerateContentResponse {
   const lastResponse = responses[responses.length - 1];
@@ -216,12 +217,19 @@ function aggregateResponses(
             ].citationMetadata!.citationSources.concat(existingMetadata);
         }
       }
-      aggregatedResponse.candidates[i].finishReason =
-        response.candidates[i].finishReason;
-      aggregatedResponse.candidates[i].finishMessage =
-        response.candidates[i].finishMessage;
-      aggregatedResponse.candidates[i].safetyRatings =
-        response.candidates[i].safetyRatings;
+      const finishResonOfChunk = response.candidates[i].finishReason;
+      if (finishResonOfChunk) {
+        aggregatedResponse.candidates[i].finishReason =
+          response.candidates[i].finishReason;
+      }
+      const finishMessageOfChunk = response.candidates[i].finishMessage;
+      if (finishMessageOfChunk) {
+        aggregatedResponse.candidates[i].finishMessage = finishMessageOfChunk;
+      }
+      const safetyRatingsOfChunk = response.candidates[i].safetyRatings;
+      if (safetyRatingsOfChunk) {
+        aggregatedResponse.candidates[i].safetyRatings = safetyRatingsOfChunk;
+      }
       if ('parts' in response.candidates[i].content) {
         for (const part of response.candidates[i].content.parts) {
           if (part.text) {

src/functions/test/post_fetch_processing_test.ts

@@ -0,0 +1,32 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {
+  AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_1,
+  STREAM_RESPONSE_CHUNKS_1,
+} from './test_data';
+import {aggregateResponses} from '../post_fetch_processing';
+
+describe('aggregateResponses', () => {
+  it('grounding metadata in muliple chunks for multiple chandidates, should aggregate accordingly', () => {
+    const actualResult = aggregateResponses(STREAM_RESPONSE_CHUNKS_1);
+
+    expect(JSON.stringify(actualResult)).toEqual(
+      JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_1)
+    );
+  });
+});