fix(js): stream tool message generated by interrupt resume (#1870)

firebase · Feb 5, 2025 · 698aedf · 698aedf
1 parent 8f1132a
commit 698aedf
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 7 deletions.
diff --git a/js/ai/src/generate/action.ts b/js/ai/src/generate/action.ts
@@ -227,10 +227,11 @@ async function generate(
   // check to make sure we don't have overlapping tool names *before* generation
   await assertValidToolNames(tools);
 
-  const { revisedRequest, interruptedResponse } = await resolveResumeOption(
-    registry,
-    rawRequest
-  );
+  const {
+    revisedRequest,
+    interruptedResponse,
+    toolMessage: resumedToolMessage,
+  } = await resolveResumeOption(registry, rawRequest);
   // NOTE: in the future we should make it possible to interrupt a restart, but
   // at the moment it's too complicated because it's not clear how to return a
   // response that amends history but doesn't generate a new message, so we throw
@@ -260,7 +261,7 @@ async function generate(
     role: Role,
     chunk: GenerateResponseChunkData
   ): GenerateResponseChunk => {
-    if (role !== chunkRole) messageIndex++;
+    if (role !== chunkRole && previousChunks.length) messageIndex++;
     chunkRole = role;
 
     const prevToSend = [...previousChunks];
@@ -275,6 +276,12 @@ async function generate(
   };
 
   const streamingCallback = getStreamingCallback(registry);
+
+  // if resolving the 'resume' option above generated a tool message, stream it.
+  if (resumedToolMessage && streamingCallback) {
+    streamingCallback(makeChunk('tool', resumedToolMessage));
+  }
+
   const response = await runWithStreamingCallback(
     registry,
     streamingCallback &&

diff --git a/js/ai/src/generate/resolve-tool-requests.ts b/js/ai/src/generate/resolve-tool-requests.ts
@@ -335,6 +335,7 @@ export async function resolveResumeOption(
 ): Promise<{
   revisedRequest?: GenerateActionOptions;
   interruptedResponse?: GenerateResponseData;
+  toolMessage?: MessageData;
 }> {
   if (!rawRequest.resume) return { revisedRequest: rawRequest }; // no-op if no resume option
   const toolMap = toToolMap(await resolveTools(registry, rawRequest.tools));
@@ -411,6 +412,7 @@ export async function resolveResumeOption(
       resume: undefined,
       messages: [...messages, toolMessage],
     },
+    toolMessage,
   });
 }
 

diff --git a/js/genkit/tests/generate_test.ts b/js/genkit/tests/generate_test.ts
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { MessageData } from '@genkit-ai/ai';
+import { GenerateResponseChunkData, MessageData } from '@genkit-ai/ai';
 import { z } from '@genkit-ai/core';
 import * as assert from 'assert';
 import { beforeEach, describe, it } from 'node:test';
@@ -120,7 +120,7 @@ describe('generate', () => {
   });
 
   describe('explicit model', () => {
-    let ai: Genkit;
+    let ai: GenkitBeta;
 
     beforeEach(() => {
       ai = genkit({});
@@ -954,5 +954,62 @@ describe('generate', () => {
         'resuming generates a tool message containing all expected responses'
       );
     });
+
+    it('streams a generated tool message when resumed', async () => {
+      pm.handleResponse = async (request, sendChunk) => {
+        sendChunk?.({
+          role: 'model',
+          index: 0,
+          content: [{ text: 'final response' }],
+        });
+        return {
+          message: { role: 'model', content: [{ text: 'final response' }] },
+        };
+      };
+
+      const chunks: GenerateResponseChunkData[] = [];
+      await ai.generate({
+        onChunk: (chunk) => chunks.push(chunk.toJSON()),
+        messages: [
+          { role: 'user', content: [{ text: 'use the doThing tool' }] },
+          {
+            role: 'model',
+            content: [
+              {
+                toolRequest: { name: 'doThing', input: {} },
+                metadata: { interrupt: true },
+              },
+            ],
+          },
+        ],
+        resume: {
+          respond: { toolResponse: { name: 'doThing', output: 'did thing' } },
+        },
+      });
+
+      assert.deepStrictEqual(chunks, [
+        {
+          content: [
+            {
+              toolResponse: {
+                name: 'doThing',
+                output: 'did thing',
+              },
+            },
+          ],
+          index: 0,
+          role: 'tool',
+        },
+        {
+          content: [
+            {
+              text: 'final response',
+            },
+          ],
+          index: 1,
+          role: 'model',
+        },
+      ]);
+    });
   });
 });