Fix case where tool use does not have input while streaming (#2226)

aws-amplify · Nov 14, 2024 · bc6dc69 · bc6dc69
1 parent 9fd0642
commit bc6dc69
Show file tree

Hide file tree

Showing 3 changed files with 97 additions and 5 deletions.
diff --git a/.changeset/clean-hounds-exercise.md b/.changeset/clean-hounds-exercise.md
@@ -0,0 +1,5 @@
+---
+'@aws-amplify/ai-constructs': patch
+---
+
+Fix case where tool use does not have input while streaming
diff --git a/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts b/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts
@@ -823,6 +823,88 @@ void describe('Bedrock converse adapter', () => {
     });
   });
 
+  void it('handles tool use with empty input when streaming', async () => {
+    const toolOutput: ToolResultContentBlock = {
+      text: 'additionalToolOutput',
+    };
+    const toolExecuteMock = mock.fn<
+      (input: unknown) => Promise<ToolResultContentBlock>
+    >(() => Promise.resolve(toolOutput));
+    const tool: ExecutableTool = {
+      name: 'toolId',
+      description: 'tool description',
+      inputSchema: {
+        json: {},
+      },
+      execute: toolExecuteMock,
+    };
+
+    const event: ConversationTurnEvent = {
+      ...commonEvent,
+    };
+
+    const bedrockClient = new BedrockRuntimeClient();
+    const bedrockResponseQueue: Array<
+      ConverseCommandOutput | ConverseStreamCommandOutput
+    > = [];
+    const toolUse1 = {
+      toolUseId: randomUUID().toString(),
+      name: tool.name,
+      input: undefined,
+    };
+    const toolUse2 = {
+      toolUseId: randomUUID().toString(),
+      name: tool.name,
+      input: '',
+    };
+    const toolUseBedrockResponse = mockBedrockResponse(
+      [
+        {
+          toolUse: toolUse1,
+        },
+        {
+          toolUse: toolUse2,
+        },
+      ],
+      true
+    );
+    bedrockResponseQueue.push(toolUseBedrockResponse);
+    const content = [
+      {
+        text: 'finalResponse',
+      },
+    ];
+    const finalBedrockResponse = mockBedrockResponse(content, true);
+    bedrockResponseQueue.push(finalBedrockResponse);
+
+    mock.method(bedrockClient, 'send', () =>
+      Promise.resolve(bedrockResponseQueue.shift())
+    );
+
+    const adapter = new BedrockConverseAdapter(
+      event,
+      [tool],
+      bedrockClient,
+      undefined,
+      messageHistoryRetriever
+    );
+
+    const chunks: Array<StreamingResponseChunk> = await askBedrockWithStreaming(
+      adapter
+    );
+    const responseText = chunks.reduce((acc, next) => {
+      if (next.contentBlockText) {
+        acc += next.contentBlockText;
+      }
+      return acc;
+    }, '');
+    assert.strictEqual(responseText, 'finalResponse');
+
+    assert.strictEqual(toolExecuteMock.mock.calls.length, 2);
+    assert.strictEqual(toolExecuteMock.mock.calls[0].arguments[0], undefined);
+    assert.strictEqual(toolExecuteMock.mock.calls[1].arguments[0], undefined);
+  });
+
   void it('throws if tool is duplicated', () => {
     assert.throws(
       () =>
@@ -1007,19 +1089,21 @@ const mockConverseStreamCommandOutput = (
           },
         },
       });
-      const input = JSON.stringify(block.toolUse.input);
+      const input = block.toolUse.input
+        ? JSON.stringify(block.toolUse.input)
+        : undefined;
       streamItems.push({
         contentBlockDelta: {
           contentBlockIndex: i,
           delta: {
             toolUse: {
               // simulate chunked input
-              input: input.substring(0, 1),
+              input: input?.substring(0, 1),
             },
           },
         },
       });
-      if (input.length > 1) {
+      if (input && input.length > 1) {
         streamItems.push({
           contentBlockDelta: {
             contentBlockIndex: i,

diff --git a/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts b/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts
@@ -232,8 +232,9 @@ export class BedrockConverseAdapter {
           if (chunk.contentBlockDelta.delta?.toolUse) {
             if (!chunk.contentBlockDelta.delta.toolUse.input) {
               toolUseInput = '';
+            } else {
+              toolUseInput += chunk.contentBlockDelta.delta.toolUse.input;
             }
-            toolUseInput += chunk.contentBlockDelta.delta.toolUse.input;
           } else if (chunk.contentBlockDelta.delta?.text) {
             text += chunk.contentBlockDelta.delta.text;
             yield {
@@ -249,7 +250,9 @@ export class BedrockConverseAdapter {
           }
         } else if (chunk.contentBlockStop) {
           if (toolUseBlock) {
-            toolUseBlock.toolUse.input = JSON.parse(toolUseInput);
+            if (toolUseInput) {
+              toolUseBlock.toolUse.input = JSON.parse(toolUseInput);
+            }
             accumulatedAssistantMessage.content?.push(toolUseBlock);
             if (
               toolUseBlock.toolUse.name &&