feat(ai): Add method to send function responses in a live session

Author: dlarocque

.changeset/lazy-donuts-agree.md

@@ -0,0 +1,8 @@
+---
+'@firebase/ai': minor
+---
+
+Added a `sendFunctionResponses` method to `LiveSession`, allowing function responses to be sent during realtime sessions.
+Fixes an issue where function responses during audio conversations caused the WebSocket connection to close. See [GitHub Issue #9264](https://github.com/firebase/firebase-js-sdk/issues/9264).
+ - **Breaking Change**: `StartAudioConversationOptions`' `functionCallingHandler` must now return a `Promise<FunctionResponse>`.
+   This breaking change is allowed since the Live API is in Public Preview.

common/api-review/ai.api.md

@@ -988,6 +988,7 @@ export class LiveSession {
     isClosed: boolean;
     receive(): AsyncGenerator<LiveServerContent | LiveServerToolCall | LiveServerToolCallCancellation>;
     send(request: string | Array<string | Part>, turnComplete?: boolean): Promise<void>;
+    sendFunctionResponses(functionResponses: FunctionResponse[]): Promise<void>;
     sendMediaChunks(mediaChunks: GenerativeContentBlob[]): Promise<void>;
     sendMediaStream(mediaChunkStream: ReadableStream<GenerativeContentBlob>): Promise<void>;
     }
@@ -1254,7 +1255,7 @@ export function startAudioConversation(liveSession: LiveSession, options?: Start
 
 // @beta
 export interface StartAudioConversationOptions {
-    functionCallingHandler?: (functionCalls: LiveServerToolCall['functionCalls']) => Promise<Part>;
+    functionCallingHandler?: (functionCalls: LiveServerToolCall['functionCalls']) => Promise<FunctionResponse>;
 }
 
 // @public

packages/ai/src/methods/live-session-helpers.test.ts

@@ -21,7 +21,11 @@ import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import { AIError } from '../errors';
 import { startAudioConversation } from './live-session-helpers';
-import { LiveServerContent, LiveServerToolCall, Part } from '../types';
+import {
+  FunctionResponse,
+  LiveServerContent,
+  LiveServerToolCall
+} from '../types';
 import { logger } from '../logger';
 import { isNode } from '@firebase/util';
 
@@ -62,6 +66,7 @@ class MockLiveSession {
   inConversation = false;
   send = sinon.stub();
   sendMediaChunks = sinon.stub();
+  sendFunctionResponses = sinon.stub();
   messageGenerator = new MockMessageGenerator();
   receive = (): MockMessageGenerator => this.messageGenerator;
 }
@@ -249,16 +254,21 @@ describe('Audio Conversation Helpers', () => {
     });
 
     it('should call function handler and send result on toolCall message.', async () => {
-      const handlerStub = sinon.stub().resolves({
-        functionResponse: { name: 'get_weather', response: { temp: '72F' } }
-      } as Part);
+      const functionResponse: FunctionResponse = {
+        id: '1',
+        name: 'get_weather',
+        response: { temp: '72F' }
+      };
+      const handlerStub = sinon.stub().resolves(functionResponse);
       const controller = await startAudioConversation(liveSession as any, {
         functionCallingHandler: handlerStub
       });
 
       const toolCallMessage: LiveServerToolCall = {
         type: 'toolCall',
-        functionCalls: [{ name: 'get_weather', args: { location: 'LA' } }]
+        functionCalls: [
+          { id: '1', name: 'get_weather', args: { location: 'LA' } }
+        ]
       };
 
       liveSession.messageGenerator.simulateMessage(toolCallMessage);
@@ -267,8 +277,8 @@ describe('Audio Conversation Helpers', () => {
       expect(handlerStub).to.have.been.calledOnceWith(
         toolCallMessage.functionCalls
       );
-      expect(liveSession.send).to.have.been.calledOnceWith([
-        { functionResponse: { name: 'get_weather', response: { temp: '72F' } } }
+      expect(liveSession.sendFunctionResponses).to.have.been.calledOnceWith([
+        functionResponse
       ]);
       await controller.stop();
     });

packages/ai/src/methods/live-session-helpers.ts

@@ -19,10 +19,10 @@ import { AIError } from '../errors';
 import { logger } from '../logger';
 import {
   AIErrorCode,
+  FunctionCall,
+  FunctionResponse,
   GenerativeContentBlob,
-  LiveServerContent,
-  LiveServerToolCall,
-  Part
+  LiveServerContent
 } from '../types';
 import { LiveSession } from './live-session';
 import { Deferred } from '@firebase/util';
@@ -115,8 +115,8 @@ export interface StartAudioConversationOptions {
    * which will then be sent back to the model.
    */
   functionCallingHandler?: (
-    functionCalls: LiveServerToolCall['functionCalls']
-  ) => Promise<Part>;
+    functionCalls: FunctionCall[]
+  ) => Promise<FunctionResponse>;
 }
 
 /**
@@ -338,11 +338,11 @@ export class AudioConversationRunner {
           );
         } else {
           try {
-            const resultPart = await this.options.functionCallingHandler(
+            const functionResponse = await this.options.functionCallingHandler(
               message.functionCalls
             );
             if (!this.isStopped) {
-              void this.liveSession.send([resultPart]);
+              void this.liveSession.sendFunctionResponses([functionResponse]);
             }
           } catch (e) {
             throw new AIError(

packages/ai/src/methods/live-session.test.ts

@@ -20,6 +20,7 @@ import { spy, stub } from 'sinon';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import {
+  FunctionResponse,
   LiveResponseType,
   LiveServerContent,
   LiveServerToolCall,
@@ -153,6 +154,35 @@ describe('LiveSession', () => {
     });
   });
 
+  describe('sendFunctionResponses()', () => {
+    it('should send all function responses', async () => {
+      const functionResponses: FunctionResponse[] = [
+        {
+          id: 'function-call-1',
+          name: 'function-name',
+          response: {
+            result: 'foo'
+          }
+        },
+        {
+          id: 'function-call-2',
+          name: 'function-name-2',
+          response: {
+            result: 'bar'
+          }
+        }
+      ];
+      await session.sendFunctionResponses(functionResponses);
+      expect(mockHandler.send).to.have.been.calledOnce;
+      const sentData = JSON.parse(mockHandler.send.getCall(0).args[0]);
+      expect(sentData).to.deep.equal({
+        toolResponse: {
+          functionResponses
+        }
+      });
+    });
+  });
+
   describe('receive()', () => {
     it('should correctly parse and transform all server message types', async () => {
       const receivePromise = (async () => {

packages/ai/src/methods/live-session.ts

@@ -17,6 +17,7 @@
 
 import {
   AIErrorCode,
+  FunctionResponse,
   GenerativeContentBlob,
   LiveResponseType,
   LiveServerContent,
@@ -30,7 +31,8 @@ import { WebSocketHandler } from '../websocket';
 import { logger } from '../logger';
 import {
   _LiveClientContent,
-  _LiveClientRealtimeInput
+  _LiveClientRealtimeInput,
+  _LiveClientToolResponse
 } from '../types/live-responses';
 
 /**
@@ -119,6 +121,32 @@ export class LiveSession {
     });
   }
 
+  /**
+   * Sends function responses to the server.
+   *
+   * @param functionResponses - The function responses to send.
+   * @throw If this session has been closed.
+   *
+   * @beta
+   */
+  async sendFunctionResponses(
+    functionResponses: FunctionResponse[]
+  ): Promise<void> {
+    if (this.isClosed) {
+      throw new AIError(
+        AIErrorCode.REQUEST_ERROR,
+        'This LiveSession has been closed and cannot be used.'
+      );
+    }
+
+    const message: _LiveClientToolResponse = {
+      toolResponse: {
+        functionResponses
+      }
+    };
+    this.webSocketHandler.send(JSON.stringify(message));
+  }
+
   /**
    * Sends a stream of {@link GenerativeContentBlob}.
    *

packages/ai/src/types/live-responses.ts

@@ -15,7 +15,12 @@
  * limitations under the License.
  */
 
-import { Content, GenerativeContentBlob, Part } from './content';
+import {
+  Content,
+  FunctionResponse,
+  GenerativeContentBlob,
+  Part
+} from './content';
 import { LiveGenerationConfig, Tool, ToolConfig } from './requests';
 
 /**
@@ -42,6 +47,17 @@ export interface _LiveClientRealtimeInput {
     mediaChunks: GenerativeContentBlob[];
   };
 }
+
+/**
+ * Function responses that is sent to the model in real time.
+ */
+// eslint-disable-next-line @typescript-eslint/naming-convention
+export interface _LiveClientToolResponse {
+  toolResponse: {
+    functionResponses: FunctionResponse[];
+  };
+}
+
 /**
  * The first message in a Live session, used to configure generation options.
  *