feat: add function calling support

PiperOrigin-RevId: 599199680

Author: sararob

README.md

@@ -167,6 +167,98 @@ async function countTokens() {
 countTokens();
 ```
 
+## Function calling
+
+The Node SDK supports
+[function calling](https://cloud.google.com/vertex-ai/docs/generative-ai/multimodal/function-calling) via `sendMessage`, `sendMessageStream`, `generateContent`, and `generateContentStream`. We recommend using it through chat methods
+(`sendMessage` or `sendMessageStream`) but have included examples of both
+approaches below.
+
+### Function declarations and response
+
+This is an example of a function declaration and function response, which are
+passed to the model in the snippets that follow.
+
+```typescript
+const functionDeclarations = [
+  {
+    function_declarations: [
+      {
+        name: "get_current_weather",
+        description: 'get weather in a given location',
+        parameters: {
+          type: FunctionDeclarationSchemaType.OBJECT,
+          properties: {
+            location: {type: FunctionDeclarationSchemaType.STRING},
+            unit: {
+              type: FunctionDeclarationSchemaType.STRING,
+              enum: ['celsius', 'fahrenheit'],
+            },
+          },
+          required: ['location'],
+        },
+      },
+    ],
+  },
+];
+
+const functionResponseParts = [
+  {
+    functionResponse: {
+      name: "get_current_weather",
+      response:
+          {name: "get_current_weather", content: {weather: "super nice"}},
+    },
+  },
+];
+```
+
+### Function calling with chat
+
+```typescript
+// Create a chat session and pass your function declarations
+const chat = generativeModel.startChat({
+  tools: functionDeclarations,
+});
+
+const chatInput1 = 'What is the weather in Boston?';
+
+// This should include a functionCall response from the model
+const result1 = await chat.sendMessageStream(chatInput1);
+for await (const item of result1.stream) {
+  console.log(item.candidates[0]);
+}
+const response1 = await result1.response;
+
+// Send a follow up message with a FunctionResponse
+const result2 = await chat.sendMessageStream(functionResponseParts);
+for await (const item of result2.stream) {
+  console.log(item.candidates[0]);
+}
+
+// This should include a text response from the model using the response content
+// provided above
+const response2 = await result2.response;
+```
+
+### Function calling with generateContentStream
+
+```typescript
+const request = {
+  contents: [
+    {role: 'user', parts: [{text: 'What is the weather in Boston?'}]},
+    {role: 'model', parts: [{functionCall: {name: 'get_current_weather', args: {'location': 'Boston'}}}]},
+    {role: 'function', parts: functionResponseParts}
+  ],
+  tools: functionDeclarations,
+};
+const streamingResp =
+    await generativeModel.generateContentStream(request);
+for await (const item of streamingResp.stream) {
+  console.log(item.candidates[0]);
+}
+```
+
 ## License
 
 The contents of this repository are licensed under the

src/index.ts

@@ -34,6 +34,7 @@ import {
   Part,
   SafetySetting,
   StreamGenerateContentResult,
+  Tool,
   VertexInit,
 } from './types/content';
 import {
@@ -134,7 +135,8 @@ export class VertexAI_Preview {
       this,
       modelParams.model,
       modelParams.generation_config,
-      modelParams.safety_settings
+      modelParams.safety_settings,
+      modelParams.tools
     );
   }
 
@@ -185,6 +187,7 @@ export declare interface StartChatParams {
   history?: Content[];
   safety_settings?: SafetySetting[];
   generation_config?: GenerationConfig;
+  tools?: Tool[];
 }
 
 // StartChatSessionRequest and ChatSession are defined here instead of in
@@ -216,6 +219,7 @@ export class ChatSession {
   private _send_stream_promise: Promise<void> = Promise.resolve();
   generation_config?: GenerationConfig;
   safety_settings?: SafetySetting[];
+  tools?: Tool[];
 
   get history(): Content[] {
     return this.historyInternal;
@@ -231,6 +235,9 @@ export class ChatSession {
     this._model_instance = request._model_instance;
     this.historyInternal = request.history ?? [];
     this._vertex_instance = request._vertex_instance;
+    this.generation_config = request.generation_config;
+    this.safety_settings = request.safety_settings;
+    this.tools = request.tools;
   }
 
   /**
@@ -241,11 +248,12 @@ export class ChatSession {
   async sendMessage(
     request: string | Array<string | Part>
   ): Promise<GenerateContentResult> {
-    const newContent: Content = formulateNewContent(request);
+    const newContent: Content[] = formulateNewContent(request);
     const generateContentrequest: GenerateContentRequest = {
-      contents: this.historyInternal.concat([newContent]),
+      contents: this.historyInternal.concat(newContent),
       safety_settings: this.safety_settings,
       generation_config: this.generation_config,
+      tools: this.tools,
     };
 
     const generateContentResult: GenerateContentResult =
@@ -257,7 +265,7 @@ export class ChatSession {
     const generateContentResponse = generateContentResult.response;
     // Only push the latest message to history if the response returned a result
     if (generateContentResponse.candidates.length !== 0) {
-      this.historyInternal.push(newContent);
+      this.historyInternal = this.historyInternal.concat(newContent);
       const contentFromAssistant =
         generateContentResponse.candidates[0].content;
       if (!contentFromAssistant.role) {
@@ -274,15 +282,15 @@ export class ChatSession {
 
   async appendHistory(
     streamGenerateContentResultPromise: Promise<StreamGenerateContentResult>,
-    newContent: Content
+    newContent: Content[]
   ): Promise<void> {
     const streamGenerateContentResult =
       await streamGenerateContentResultPromise;
     const streamGenerateContentResponse =
       await streamGenerateContentResult.response;
     // Only push the latest message to history if the response returned a result
     if (streamGenerateContentResponse.candidates.length !== 0) {
-      this.historyInternal.push(newContent);
+      this.historyInternal = this.historyInternal.concat(newContent);
       const contentFromAssistant =
         streamGenerateContentResponse.candidates[0].content;
       if (!contentFromAssistant.role) {
@@ -303,11 +311,12 @@ export class ChatSession {
   async sendMessageStream(
     request: string | Array<string | Part>
   ): Promise<StreamGenerateContentResult> {
-    const newContent: Content = formulateNewContent(request);
+    const newContent: Content[] = formulateNewContent(request);
     const generateContentrequest: GenerateContentRequest = {
-      contents: this.historyInternal.concat([newContent]),
+      contents: this.historyInternal.concat(newContent),
       safety_settings: this.safety_settings,
       generation_config: this.generation_config,
+      tools: this.tools,
     };
 
     const streamGenerateContentResultPromise = this._model_instance
@@ -335,6 +344,7 @@ export class GenerativeModel {
   model: string;
   generation_config?: GenerationConfig;
   safety_settings?: SafetySetting[];
+  tools?: Tool[];
   private _vertex_instance: VertexAI_Preview;
   private _use_non_stream = false;
   private publisherModelEndpoint: string;
@@ -351,12 +361,14 @@ export class GenerativeModel {
     vertex_instance: VertexAI_Preview,
     model: string,
     generation_config?: GenerationConfig,
-    safety_settings?: SafetySetting[]
+    safety_settings?: SafetySetting[],
+    tools?: Tool[]
   ) {
     this._vertex_instance = vertex_instance;
     this.model = model;
     this.generation_config = generation_config;
     this.safety_settings = safety_settings;
+    this.tools = tools;
     if (model.startsWith('models/')) {
       this.publisherModelEndpoint = `publishers/google/${this.model}`;
     } else {
@@ -401,6 +413,7 @@ export class GenerativeModel {
       contents: request.contents,
       generation_config: request.generation_config ?? this.generation_config,
       safety_settings: request.safety_settings ?? this.safety_settings,
+      tools: request.tools ?? [],
     };
 
     const response: Response | undefined = await postRequest({
@@ -444,6 +457,7 @@ export class GenerativeModel {
       contents: request.contents,
       generation_config: request.generation_config ?? this.generation_config,
       safety_settings: request.safety_settings ?? this.safety_settings,
+      tools: request.tools ?? [],
     };
     const response = await postRequest({
       region: this._vertex_instance.location,
@@ -501,12 +515,15 @@ export class GenerativeModel {
         request.generation_config ?? this.generation_config;
       startChatRequest.safety_settings =
         request.safety_settings ?? this.safety_settings;
+      startChatRequest.tools = request.tools ?? this.tools;
     }
     return new ChatSession(startChatRequest);
   }
 }
 
-function formulateNewContent(request: string | Array<string | Part>): Content {
+function formulateNewContent(
+  request: string | Array<string | Part>
+): Content[] {
   let newParts: Part[] = [];
 
   if (typeof request === 'string') {
@@ -521,8 +538,38 @@ function formulateNewContent(request: string | Array<string | Part>): Content {
     }
   }
 
-  const newContent: Content = {role: constants.USER_ROLE, parts: newParts};
-  return newContent;
+  return formatPartsByRole(newParts);
+}
+
+/**
+ * When multiple Part types (i.e. FunctionResponsePart and TextPart) are
+ * passed in a single Part array, we may need to assign different roles to each
+ * part. Currently only FunctionResponsePart requires a role other than 'user'.
+ * @ignore
+ * @param {Array<Part>} parts Array of parts to pass to the model
+ * @return {Content[]} Array of content items
+ */
+function formatPartsByRole(parts: Array<Part>): Content[] {
+  const partsByRole: Content[] = [];
+  const userContent: Content = {role: constants.USER_ROLE, parts: []};
+  const functionContent: Content = {role: constants.FUNCTION_ROLE, parts: []};
+
+  for (const part of parts) {
+    if ('functionResponse' in part) {
+      functionContent.parts.push(part);
+    } else {
+      userContent.parts.push(part);
+    }
+  }
+
+  if (userContent.parts.length > 0) {
+    partsByRole.push(userContent);
+  }
+  if (functionContent.parts.length > 0) {
+    partsByRole.push(functionContent);
+  }
+
+  return partsByRole;
 }
 
 function throwErrorIfNotOK(response: Response | undefined) {

src/process_stream.ts

@@ -193,6 +193,13 @@ function aggregateResponses(
           if (part.text) {
             aggregatedResponse.candidates[i].content.parts[0].text += part.text;
           }
+          if (part.functionCall) {
+            aggregatedResponse.candidates[i].content.parts[0].functionCall =
+              part.functionCall;
+            // the empty 'text' key should be removed if functionCall is in the
+            // response
+            delete aggregatedResponse.candidates[i].content.parts[0].text;
+          }
         }
       }
     }