chore: Move functionCalls quick accessor in Candidate interface to a seperate helper function

PiperOrigin-RevId: 618039934

Author: qiaodev

src/functions/index.ts

@@ -18,3 +18,4 @@
 export {countTokens} from './count_tokens';
 export {postRequest} from './post_request';
 export {generateContent, generateContentStream} from './generate_content';
+export {ResponseHandler} from './response_handler';

src/functions/post_fetch_processing.ts

@@ -58,7 +58,7 @@ async function* generateResponseSequence(
     if (done) {
       break;
     }
-    yield addMissingFields(value);
+    yield addMissingIndexAndRole(value);
   }
 }
 
@@ -340,37 +340,6 @@ function addMissingIndexAndRole(
   return generateContentResponse;
 }
 
-function addCandidateFunctionCalls(
-  response: GenerateContentResponse
-): GenerateContentResponse {
-  if (!response.candidates) {
-    return response;
-  }
-  for (const candidate of response.candidates) {
-    if (
-      !candidate.content ||
-      !candidate.content.parts ||
-      candidate.content.parts.length === 0
-    ) {
-      continue;
-    }
-    const functionCalls = candidate.content.parts
-      .filter((part: Part) => !!part.functionCall)
-      .map((part: Part) => part.functionCall!);
-    if (functionCalls.length > 0) {
-      candidate.functionCalls = functionCalls;
-    }
-  }
-  return response;
-}
-
-function addMissingFields(
-  response: GenerateContentResponse
-): GenerateContentResponse {
-  const generateContentResponse = addMissingIndexAndRole(response);
-  return addCandidateFunctionCalls(generateContentResponse);
-}
-
 /**
  * Process model responses from generateContent
  * @ignore
@@ -383,7 +352,7 @@ export async function processUnary(
     const responseJson = await response.json();
     const generateContentResponse = addMissingIndexAndRole(responseJson);
     return Promise.resolve({
-      response: addCandidateFunctionCalls(generateContentResponse),
+      response: generateContentResponse,
     });
   }
 

src/functions/response_handler.ts

@@ -0,0 +1,36 @@
+/**
+ * @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 {GenerateContentCandidate, FunctionCall, Part} from '../types';
+
+/** Helper class to render any extra properties out of {@link GenerateContentResponse}. */
+export class ResponseHandler {
+  /**
+   * Extracts function calls from a {@link GenerateContentCandidate}.
+   *
+   * @param candidate - The candidate to extract function calls from.
+   * @returns the array of function calls in a {@link GenerateContentCandidate}.
+   */
+  static getCandidateFunctionCalls(
+    candidate?: GenerateContentCandidate
+  ): FunctionCall[] {
+    if (!candidate) return [] as FunctionCall[];
+    return candidate.content.parts
+      .filter((part: Part | undefined) => !!part && !!part.functionCall)
+      .map((part: Part) => part.functionCall!);
+  }
+}

src/functions/test/functions_test.ts

@@ -36,6 +36,7 @@ import {constants} from '../../util';
 import {countTokens} from '../count_tokens';
 import {generateContent, generateContentStream} from '../generate_content';
 import * as StreamFunctions from '../post_fetch_processing';
+import {ResponseHandler} from '../response_handler';
 
 const TEST_PROJECT = 'test-project';
 const TEST_LOCATION = 'test-location';
@@ -150,7 +151,6 @@ const TEST_CANDIDATES_WITH_FUNCTION_CALL = [
     finishReason: FinishReason.STOP,
     finishMessage: '',
     safetyRatings: TEST_SAFETY_RATINGS,
-    functionCalls: [TEST_FUNCTION_CALL_RESPONSE.functionCall],
   },
 ];
 const TEST_MODEL_RESPONSE_WITH_FUNCTION_CALL = {
@@ -617,8 +617,16 @@ describe('generateContent', () => {
       TEST_API_ENDPOINT
     );
     expect(actualResult).toEqual(expectedResult);
-    expect(actualResult.response.candidates![0].functionCalls).toHaveSize(1);
-    expect(actualResult.response.candidates![0].functionCalls).toEqual([
+    expect(
+      ResponseHandler.getCandidateFunctionCalls(
+        actualResult.response.candidates![0]
+      )
+    ).toHaveSize(1);
+    expect(
+      ResponseHandler.getCandidateFunctionCalls(
+        actualResult.response.candidates![0]
+      )
+    ).toEqual([
       expectedResult.response.candidates![0].content.parts[0].functionCall!,
     ]);
   });
@@ -653,8 +661,10 @@ describe('generateContent', () => {
     );
     expect(actualResult).toEqual(expectedResult);
     expect(
-      actualResult.response.candidates![0].functionCalls
-    ).not.toBeDefined();
+      ResponseHandler.getCandidateFunctionCalls(
+        actualResult.response.candidates?.[0]
+      )
+    ).toHaveSize(0);
   });
 
   it('returns empty candidates when response is empty', async () => {

src/index.ts

@@ -18,3 +18,4 @@
 export {VertexAI} from './vertex_ai';
 export * from './types';
 export * from './models';
+export {ResponseHandler} from './functions';

src/types/content.ts

@@ -525,8 +525,6 @@ export declare interface GenerateContentCandidate {
   citationMetadata?: CitationMetadata;
   /** Optional. {@link GroundingMetadata}. */
   groundingMetadata?: GroundingMetadata;
-  /* Optional. Array of {@link FunctionCall}. */
-  functionCalls?: FunctionCall[];
 }
 
 /**

system_test/end_to_end_sample_test.ts

@@ -23,6 +23,7 @@ import {
   Part,
   TextPart,
   VertexAI,
+  ResponseHandler,
 } from '../src';
 import {FunctionDeclarationSchemaType} from '../src/types';
 
@@ -472,7 +473,9 @@ describe('generateContentStream', () => {
         .candidates![0].content.parts.filter(part => !!part.functionCall)
         .map(part => part.functionCall!);
       expect(functionCalls).toHaveSize(1);
-      expect(item.candidates![0].functionCalls!).toEqual(functionCalls!);
+      expect(
+        ResponseHandler.getCandidateFunctionCalls(item.candidates?.[0])
+      ).toEqual(functionCalls!);
     }
   });
   it('in preview should return a FunctionCall when passed a FunctionDeclaration', async () => {
@@ -494,7 +497,9 @@ describe('generateContentStream', () => {
         .candidates![0].content.parts.filter(part => !!part.functionCall)
         .map(part => part.functionCall!);
       expect(functionCalls).toHaveSize(1);
-      expect(item.candidates![0].functionCalls!).toEqual(functionCalls!);
+      expect(
+        ResponseHandler.getCandidateFunctionCalls(item.candidates?.[0])
+      ).toEqual(functionCalls!);
     }
   });
   xit('should return grounding metadata when passed GoogleSearchRetriever in getGenerativeModel', async () => {
@@ -770,7 +775,9 @@ describe('generateContent', () => {
       .candidates![0].content.parts.filter((part: Part) => !!part.functionCall)
       .map((part: Part) => part.functionCall!);
     expect(functionCalls).toHaveSize(1);
-    expect(resp.response.candidates![0].functionCalls!).toEqual(functionCalls!);
+    expect(
+      ResponseHandler.getCandidateFunctionCalls(resp.response.candidates![0])
+    ).toEqual(functionCalls!);
   });
   it('in preview should return a FunctionCall when passed a FunctionDeclaration', async () => {
     const request = {
@@ -789,7 +796,9 @@ describe('generateContent', () => {
       .candidates![0].content.parts.filter((part: Part) => !!part.functionCall)
       .map((part: Part) => part.functionCall!);
     expect(functionCalls).toHaveSize(1);
-    expect(resp.response.candidates![0].functionCalls!).toEqual(functionCalls!);
+    expect(
+      ResponseHandler.getCandidateFunctionCalls(resp.response.candidates![0])
+    ).toEqual(functionCalls!);
   });
 });