feat(ai): support generateContent function auto-calling

Author: mikehardy

packages/ai/__tests__/generative-model.test.ts

@@ -21,6 +21,7 @@ import { AI, FunctionCallingMode, ThinkingLevel } from '../lib/public-types';
 import * as request from '../lib/requests/request';
 import { BackendName, getMockResponse } from './test-utils/mock-response';
 import { VertexAIBackend } from '../lib/backend';
+import { GenerateContentResponse } from '../lib';
 
 const fakeAI: AI = {
   app: {
@@ -36,6 +37,10 @@ const fakeAI: AI = {
   location: 'us-central1',
 };
 
+function responseFromJson(json: GenerateContentResponse): Response {
+  return { json: async () => json } as Response;
+}
+
 describe('GenerativeModel', () => {
   it('passes CodeExecutionTool and URLContextTool with other tools through to generateContent', async function () {
     const genModel = new GenerativeModel(fakeAI, {
@@ -362,6 +367,151 @@ describe('GenerativeModel', () => {
     makeRequestStub.mockRestore();
   });
 
+  it('automatically calls functionReference from generateContent function calls', async () => {
+    const getWeather = jest.fn<(args: object) => object>().mockReturnValue({ temperature: 72 });
+    const genModel = new GenerativeModel(fakeAI, {
+      model: 'my-model',
+      tools: [
+        {
+          functionDeclarations: [
+            {
+              name: 'getWeather',
+              description: 'Gets weather for a city.',
+              functionReference: getWeather,
+            },
+          ],
+        },
+      ],
+    });
+    const makeRequestStub = jest
+      .spyOn(request, 'makeRequest')
+      .mockResolvedValueOnce(
+        responseFromJson({
+          candidates: [
+            {
+              index: 0,
+              content: {
+                role: 'model',
+                parts: [
+                  {
+                    functionCall: {
+                      id: 'call-1',
+                      name: 'getWeather',
+                      args: { city: 'London' },
+                    },
+                  },
+                ],
+              },
+            },
+          ],
+        }),
+      )
+      .mockResolvedValueOnce(
+        responseFromJson({
+          candidates: [
+            {
+              index: 0,
+              content: {
+                role: 'model',
+                parts: [{ text: 'It is 72 degrees.' }],
+              },
+            },
+          ],
+        }),
+      );
+
+    const result = await genModel.generateContent('weather in London');
+
+    expect(result.response.text()).toBe('It is 72 degrees.');
+    expect(getWeather).toHaveBeenCalledWith({ city: 'London' });
+    expect(makeRequestStub).toHaveBeenCalledTimes(2);
+    const followUpBody = JSON.parse(makeRequestStub.mock.calls[1]![1] as string);
+    expect(followUpBody.contents).toEqual([
+      {
+        role: 'user',
+        parts: [{ text: 'weather in London' }],
+      },
+      {
+        role: 'model',
+        parts: [
+          {
+            functionCall: {
+              id: 'call-1',
+              name: 'getWeather',
+              args: { city: 'London' },
+            },
+          },
+        ],
+      },
+      {
+        role: 'function',
+        parts: [
+          {
+            functionResponse: {
+              id: 'call-1',
+              name: 'getWeather',
+              response: { temperature: 72 },
+            },
+          },
+        ],
+      },
+    ]);
+    makeRequestStub.mockRestore();
+  });
+
+  it('returns the latest response when maxSequentialFunctionCalls is reached', async () => {
+    const getWeather = jest.fn<(args: object) => object>().mockReturnValue({ temperature: 72 });
+    const genModel = new GenerativeModel(
+      fakeAI,
+      {
+        model: 'my-model',
+        tools: [
+          {
+            functionDeclarations: [
+              {
+                name: 'getWeather',
+                description: 'Gets weather for a city.',
+                functionReference: getWeather,
+              },
+            ],
+          },
+        ],
+      },
+      { maxSequentialFunctionCalls: 1 },
+    );
+    const functionCallResponse: GenerateContentResponse = {
+      candidates: [
+        {
+          index: 0,
+          content: {
+            role: 'model',
+            parts: [
+              {
+                functionCall: {
+                  name: 'getWeather',
+                  args: { city: 'London' },
+                },
+              },
+            ],
+          },
+        },
+      ],
+    };
+    const makeRequestStub = jest
+      .spyOn(request, 'makeRequest')
+      .mockResolvedValueOnce(responseFromJson(functionCallResponse))
+      .mockResolvedValueOnce(responseFromJson(functionCallResponse));
+
+    const result = await genModel.generateContent('weather in London');
+
+    expect(result.response.functionCalls()).toEqual([
+      { name: 'getWeather', args: { city: 'London' } },
+    ]);
+    expect(getWeather).toHaveBeenCalledTimes(1);
+    expect(makeRequestStub).toHaveBeenCalledTimes(2);
+    makeRequestStub.mockRestore();
+  });
+
   it('passes base model params through to ChatSession when there are no startChatParams', async () => {
     const genModel = new GenerativeModel(fakeAI, {
       model: 'my-model',

packages/ai/lib/models/generative-model.ts

@@ -20,6 +20,9 @@ import {
   Content,
   CountTokensRequest,
   CountTokensResponse,
+  FunctionCall,
+  FunctionDeclaration,
+  FunctionResponse,
   GenerateContentRequest,
   GenerateContentResult,
   GenerateContentStreamResult,
@@ -40,6 +43,8 @@ import { mergeRequestOptions } from '../requests/request-options';
 import { AIModel } from './ai-model';
 import { AI } from '../public-types';
 
+const DEFAULT_MAX_SEQUENTIAL_FUNCTION_CALLS = 10;
+
 /**
  * Class for generative model APIs.
  * @public
@@ -71,19 +76,17 @@ export class GenerativeModel extends AIModel {
     singleRequestOptions?: SingleRequestOptions,
   ): Promise<GenerateContentResult> {
     const formattedParams = formatGenerateContentInput(request);
-    return generateContent(
-      this._apiSettings,
-      this.model,
-      {
-        generationConfig: this.generationConfig,
-        safetySettings: this.safetySettings,
-        tools: this.tools,
-        toolConfig: this.toolConfig,
-        systemInstruction: this.systemInstruction,
-        ...formattedParams,
-      },
-      mergeRequestOptions(this.requestOptions, singleRequestOptions),
-    );
+    const params: GenerateContentRequest = {
+      generationConfig: this.generationConfig,
+      safetySettings: this.safetySettings,
+      tools: this.tools,
+      toolConfig: this.toolConfig,
+      systemInstruction: this.systemInstruction,
+      ...formattedParams,
+    };
+    const requestOptions = mergeRequestOptions(this.requestOptions, singleRequestOptions);
+    const result = await generateContent(this._apiSettings, this.model, params, requestOptions);
+    return this._generateContentWithAutomaticFunctionCalling(params, result, requestOptions);
   }
 
   /**
@@ -152,4 +155,92 @@ export class GenerativeModel extends AIModel {
       mergeRequestOptions(this.requestOptions, singleRequestOptions),
     );
   }
+
+  private async _generateContentWithAutomaticFunctionCalling(
+    params: GenerateContentRequest,
+    result: GenerateContentResult,
+    requestOptions?: SingleRequestOptions,
+  ): Promise<GenerateContentResult> {
+    let remainingFunctionCalls =
+      requestOptions?.maxSequentialFunctionCalls ?? DEFAULT_MAX_SEQUENTIAL_FUNCTION_CALLS;
+    let currentParams = params;
+    let currentResult = result;
+
+    while (remainingFunctionCalls > 0) {
+      const functionCalls = currentResult.response.functionCalls?.();
+      if (!functionCalls?.length) {
+        return currentResult;
+      }
+
+      const functionResponses = await this._callFunctionReferences(
+        currentParams.tools,
+        functionCalls,
+      );
+      if (!functionResponses) {
+        return currentResult;
+      }
+
+      const responseContent = currentResult.response.candidates?.[0]?.content;
+      if (!responseContent) {
+        return currentResult;
+      }
+
+      remainingFunctionCalls -= 1;
+      currentParams = {
+        ...currentParams,
+        contents: [
+          ...currentParams.contents,
+          responseContent,
+          {
+            role: 'function',
+            parts: functionResponses.map(functionResponse => ({ functionResponse })),
+          },
+        ],
+      };
+      currentResult = await generateContent(
+        this._apiSettings,
+        this.model,
+        currentParams,
+        requestOptions,
+      );
+    }
+
+    return currentResult;
+  }
+
+  private async _callFunctionReferences(
+    tools: Tool[] | undefined,
+    functionCalls: FunctionCall[],
+  ): Promise<FunctionResponse[] | undefined> {
+    const declarations = this._getFunctionDeclarationsWithReferences(tools);
+    if (!declarations.length) {
+      return undefined;
+    }
+
+    const functionResponses: FunctionResponse[] = [];
+    for (const functionCall of functionCalls) {
+      const declaration = declarations.find(candidate => candidate.name === functionCall.name);
+      if (!declaration?.functionReference) {
+        return undefined;
+      }
+
+      const response = (await declaration.functionReference(functionCall.args)) as object;
+      functionResponses.push({
+        id: functionCall.id,
+        name: functionCall.name,
+        response,
+      });
+    }
+    return functionResponses;
+  }
+
+  private _getFunctionDeclarationsWithReferences(tools: Tool[] | undefined): FunctionDeclaration[] {
+    return (
+      tools?.flatMap(tool =>
+        'functionDeclarations' in tool
+          ? (tool.functionDeclarations?.filter(declaration => declaration.functionReference) ?? [])
+          : [],
+      ) ?? []
+    );
+  }
 }