fix: correct CitationMetadata interface. refactor nested function

PiperOrigin-RevId: 616874627

Author: yyyu-google

src/functions/post_fetch_processing.ts

@@ -16,8 +16,8 @@
  */
 
 import {
-  CitationSource,
-  Content,
+  Citation,
+  CitationMetadata,
   CountTokensResponse,
   GenerateContentCandidate,
   GenerateContentResponse,
@@ -199,23 +199,14 @@ export function aggregateResponses(
           },
         } as GenerateContentCandidate;
       }
-      if (response.candidates[i].citationMetadata) {
-        if (
-          !aggregatedResponse.candidates[i].citationMetadata?.citationSources
-        ) {
-          aggregatedResponse.candidates[i].citationMetadata = {
-            citationSources: [] as CitationSource[],
-          };
-        }
-
-        const existingMetadata = response.candidates[i].citationMetadata ?? {};
-
-        if (aggregatedResponse.candidates[i].citationMetadata) {
-          aggregatedResponse.candidates[i].citationMetadata!.citationSources =
-            aggregatedResponse.candidates[
-              i
-            ].citationMetadata!.citationSources.concat(existingMetadata);
-        }
+      const citationMetadataAggregated: CitationMetadata | undefined =
+        aggregateCitationMetadataForCandidate(
+          response.candidates[i],
+          aggregatedResponse.candidates[i]
+        );
+      if (citationMetadataAggregated) {
+        aggregatedResponse.candidates[i].citationMetadata =
+          citationMetadataAggregated;
       }
       const finishResonOfChunk = response.candidates[i].finishReason;
       if (finishResonOfChunk) {
@@ -260,6 +251,29 @@ export function aggregateResponses(
   return aggregatedResponse;
 }
 
+function aggregateCitationMetadataForCandidate(
+  candidateChunk: GenerateContentCandidate,
+  aggregatedCandidate: GenerateContentCandidate
+): CitationMetadata | undefined {
+  if (!candidateChunk.citationMetadata) {
+    return;
+  }
+  const emptyCitationMetadata: CitationMetadata = {
+    citations: [],
+  };
+  const citationMetadataAggregated: CitationMetadata =
+    aggregatedCandidate.citationMetadata ?? emptyCitationMetadata;
+  const citationMetadataChunk: CitationMetadata =
+    candidateChunk.citationMetadata!;
+  if (citationMetadataChunk.citations) {
+    citationMetadataAggregated.citations =
+      citationMetadataAggregated.citations!.concat(
+        citationMetadataChunk.citations
+      );
+  }
+  return citationMetadataAggregated;
+}
+
 function aggregateGroundingMetadataForCandidate(
   candidateChunk: GenerateContentCandidate,
   aggregatedCandidate: GenerateContentCandidate

src/functions/test/functions_test.ts

@@ -104,7 +104,7 @@ const TEST_CANDIDATES = [
     finishMessage: '',
     safetyRatings: TEST_SAFETY_RATINGS,
     citationMetadata: {
-      citationSources: [
+      citations: [
         {
           startIndex: 367,
           endIndex: 491,
@@ -585,9 +585,9 @@ describe('generateContent', () => {
       TEST_API_ENDPOINT
     );
     expect(
-      resp.response.candidates[0].citationMetadata?.citationSources.length
+      resp.response.candidates[0].citationMetadata?.citations.length
     ).toEqual(
-      TEST_MODEL_RESPONSE.candidates[0].citationMetadata.citationSources.length
+      TEST_MODEL_RESPONSE.candidates[0].citationMetadata.citations.length
     );
   });
 

src/functions/test/post_fetch_processing_test.ts

@@ -17,16 +17,26 @@
 
 import {
   AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_1,
+  AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_2,
   STREAM_RESPONSE_CHUNKS_1,
+  STREAM_RESPONSE_CHUNKS_2,
 } from './test_data';
 import {aggregateResponses} from '../post_fetch_processing';
 
 describe('aggregateResponses', () => {
-  it('grounding metadata in muliple chunks for multiple chandidates, should aggregate accordingly', () => {
+  it('grounding metadata in multiple chunks for multiple candidates, should aggregate accordingly', () => {
     const actualResult = aggregateResponses(STREAM_RESPONSE_CHUNKS_1);
 
     expect(JSON.stringify(actualResult)).toEqual(
       JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_1)
     );
   });
+
+  it('citation metadata in multiple chunks for multiple candidates, should aggregate accordingly', () => {
+    const actualResult = aggregateResponses(STREAM_RESPONSE_CHUNKS_2);
+
+    expect(JSON.stringify(actualResult)).toEqual(
+      JSON.stringify(AGGREGATED_RESPONSE_STREAM_RESPONSE_CHUNKS_2)
+    );
+  });
 });