feat(ai): support streaming chat function auto-calling

mikehardy · mikehardy · commit 670a247268c2 · 2026-06-04T13:12:21.000-05:00
diff --git a/packages/ai/__tests__/chat-session.test.ts b/packages/ai/__tests__/chat-session.test.ts
@@ -17,7 +17,7 @@
 import { describe, expect, it, afterEach, jest } from '@jest/globals';
 
 import * as generateContentMethods from '../lib/methods/generate-content';
-import { GenerateContentStreamResult } from '../lib/types';
+import { EnhancedGenerateContentResponse, GenerateContentStreamResult } from '../lib/types';
 import { ChatSession } from '../lib/methods/chat-session';
 import { ApiSettings } from '../lib/types/internal';
 import { RequestOptions } from '../lib/types/requests';
@@ -35,6 +35,15 @@ const requestOptions: RequestOptions = {
   timeout: 1000,
 };
 
+function streamResult(response: EnhancedGenerateContentResponse): GenerateContentStreamResult {
+  return {
+    stream: (async function* () {
+      yield response;
+    })(),
+    response: Promise.resolve(response),
+  };
+}
+
 describe('ChatSession', () => {
   afterEach(() => {
     jest.restoreAllMocks();
@@ -129,6 +138,102 @@ describe('ChatSession', () => {
       );
     });
 
+    it('automatically calls functionReference from stream function calls', async () => {
+      const getWeather = jest.fn<(args: object) => object>().mockReturnValue({ temperature: 72 });
+      const functionCallResponse = {
+        candidates: [
+          {
+            index: 0,
+            content: {
+              role: 'model',
+              parts: [
+                {
+                  functionCall: {
+                    name: 'getWeather',
+                    args: { city: 'London' },
+                  },
+                },
+              ],
+            },
+          },
+        ],
+        functionCalls: () => [{ name: 'getWeather', args: { city: 'London' } }],
+      } as EnhancedGenerateContentResponse;
+      const finalResponse = {
+        candidates: [
+          {
+            index: 0,
+            content: {
+              role: 'model',
+              parts: [{ text: 'It is 72 degrees.' }],
+            },
+          },
+        ],
+        functionCalls: () => undefined,
+      } as EnhancedGenerateContentResponse;
+      const generateContentStreamStub = jest
+        .spyOn(generateContentMethods, 'generateContentStream')
+        .mockResolvedValueOnce(streamResult(functionCallResponse))
+        .mockResolvedValueOnce(streamResult(finalResponse));
+      const chatSession = new ChatSession(
+        fakeApiSettings,
+        'a-model',
+        {
+          tools: [
+            {
+              functionDeclarations: [
+                {
+                  name: 'getWeather',
+                  description: 'Gets weather for a city.',
+                  functionReference: getWeather,
+                },
+              ],
+            },
+          ],
+        },
+        requestOptions,
+      );
+
+      const result = await chatSession.sendMessageStream('weather in London');
+      await result.response;
+      const history = await chatSession.getHistory();
+
+      expect(getWeather).toHaveBeenCalledWith({ city: 'London' });
+      expect(generateContentStreamStub).toHaveBeenCalledTimes(2);
+      expect(history).toEqual([
+        {
+          role: 'user',
+          parts: [{ text: 'weather in London' }],
+        },
+        {
+          role: 'model',
+          parts: [
+            {
+              functionCall: {
+                name: 'getWeather',
+                args: { city: 'London' },
+              },
+            },
+          ],
+        },
+        {
+          role: 'function',
+          parts: [
+            {
+              functionResponse: {
+                name: 'getWeather',
+                response: { temperature: 72 },
+              },
+            },
+          ],
+        },
+        {
+          role: 'model',
+          parts: [{ text: 'It is 72 degrees.' }],
+        },
+      ]);
+    });
+
     it('downstream sendPromise errors should log but not throw', async () => {
       const consoleStub = jest.spyOn(console, 'error').mockImplementation(() => {});
       // make response undefined so that response.candidates errors
diff --git a/packages/ai/lib/methods/automatic-function-calling.ts b/packages/ai/lib/methods/automatic-function-calling.ts
@@ -22,11 +22,13 @@ import {
   FunctionResponse,
   GenerateContentRequest,
   GenerateContentResult,
+  GenerateContentStreamResult,
+  EnhancedGenerateContentResponse,
   SingleRequestOptions,
   Tool,
 } from '../types';
 import { ApiSettings } from '../types/internal';
-import { generateContent } from './generate-content';
+import { generateContent, generateContentStream } from './generate-content';
 
 const DEFAULT_MAX_SEQUENTIAL_FUNCTION_CALLS = 10;
 
@@ -35,6 +37,11 @@ export interface AutomaticFunctionCallingResult {
   addedContents: Content[];
 }
 
+export interface AutomaticFunctionCallingStreamResult {
+  result: GenerateContentStreamResult;
+  addedContents: Content[];
+}
+
 export async function generateContentWithAutomaticFunctionCalling(
   apiSettings: ApiSettings,
   model: string,
@@ -80,8 +87,61 @@ export async function generateContentWithAutomaticFunctionCalling(
   return { result: currentResult, addedContents };
 }
 
-function getModelResponseContent(result: GenerateContentResult): Content | undefined {
-  const responseContent = result.response.candidates?.[0]?.content;
+export async function generateContentStreamWithAutomaticFunctionCalling(
+  apiSettings: ApiSettings,
+  model: string,
+  params: GenerateContentRequest,
+  result: GenerateContentStreamResult,
+  requestOptions?: SingleRequestOptions,
+): Promise<AutomaticFunctionCallingStreamResult> {
+  if (!getFunctionDeclarationsWithReferences(params.tools).length) {
+    return { result, addedContents: [] };
+  }
+
+  let remainingFunctionCalls =
+    requestOptions?.maxSequentialFunctionCalls ?? DEFAULT_MAX_SEQUENTIAL_FUNCTION_CALLS;
+  let currentParams = params;
+  let currentResult = result;
+  const addedContents: Content[] = [];
+
+  while (remainingFunctionCalls > 0) {
+    const response = await currentResult.response;
+    const functionCalls = response.functionCalls?.();
+    if (!functionCalls?.length) {
+      return { result: currentResult, addedContents };
+    }
+
+    const functionResponses = await callFunctionReferences(currentParams.tools, functionCalls);
+    if (!functionResponses) {
+      return { result: currentResult, addedContents };
+    }
+
+    const responseContent = getModelResponseContent(response);
+    if (!responseContent) {
+      return { result: currentResult, addedContents };
+    }
+
+    remainingFunctionCalls -= 1;
+    const functionResponseContent: Content = {
+      role: 'function',
+      parts: functionResponses.map(functionResponse => ({ functionResponse })),
+    };
+    addedContents.push(responseContent, functionResponseContent);
+    currentParams = {
+      ...currentParams,
+      contents: [...currentParams.contents, responseContent, functionResponseContent],
+    };
+    currentResult = await generateContentStream(apiSettings, model, currentParams, requestOptions);
+  }
+
+  return { result: currentResult, addedContents };
+}
+
+function getModelResponseContent(
+  responseOrResult: GenerateContentResult | EnhancedGenerateContentResponse,
+): Content | undefined {
+  const response = 'response' in responseOrResult ? responseOrResult.response : responseOrResult;
+  const responseContent = response.candidates?.[0]?.content;
   if (!responseContent) {
     return undefined;
   }
diff --git a/packages/ai/lib/methods/chat-session.ts b/packages/ai/lib/methods/chat-session.ts
@@ -36,7 +36,10 @@ import {
   templateGenerateContent,
   templateGenerateContentStream,
 } from './generate-content';
-import { generateContentWithAutomaticFunctionCalling } from './automatic-function-calling';
+import {
+  generateContentStreamWithAutomaticFunctionCalling,
+  generateContentWithAutomaticFunctionCalling,
+} from './automatic-function-calling';
 import { ApiSettings } from '../types/internal';
 import { logger } from '../logger';
 import { mergeRequestOptions } from '../requests/request-options';
@@ -176,11 +179,20 @@ export class ChatSession extends ChatSessionBase<StartChatParams> {
       systemInstruction: this.params?.systemInstruction,
       contents: [...this._history, newContent],
     };
+    const requestOptions = mergeRequestOptions(this.requestOptions, singleRequestOptions);
     const streamPromise = generateContentStream(
       this._apiSettings,
       this.model,
       generateContentRequest,
-      mergeRequestOptions(this.requestOptions, singleRequestOptions),
+      requestOptions,
+    ).then(result =>
+      generateContentStreamWithAutomaticFunctionCalling(
+        this._apiSettings,
+        this.model,
+        generateContentRequest,
+        result,
+        requestOptions,
+      ),
     );
 
     // Add onto the chain.
@@ -191,10 +203,13 @@ export class ChatSession extends ChatSessionBase<StartChatParams> {
       .catch(_ignored => {
         throw new Error(SILENT_ERROR);
       })
-      .then(streamResult => streamResult.response)
-      .then((response: EnhancedGenerateContentResponse) => {
+      .then(({ result, addedContents }) =>
+        result.response.then(response => ({ response, addedContents })),
+      )
+      .then(({ response, addedContents }) => {
         if (response.candidates && response.candidates.length > 0) {
           this._history.push(newContent);
+          this._history.push(...addedContents);
           const responseContent = { ...response.candidates[0]?.content };
           // Response seems to come back without a role set.
           if (!responseContent.role) {
@@ -220,7 +235,7 @@ export class ChatSession extends ChatSessionBase<StartChatParams> {
           logger.error(e);
         }
       });
-    return streamPromise;
+    return (await streamPromise).result;
   }
 }
 
