Skip to content

Commit bc6dc69

Browse files
author
Kamil Sobol
authored
Fix case where tool use does not have input while streaming (#2226)
1 parent 9fd0642 commit bc6dc69

File tree

3 files changed

+97
-5
lines changed

3 files changed

+97
-5
lines changed
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@aws-amplify/ai-constructs': patch
3+
---
4+
5+
Fix case where tool use does not have input while streaming

packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts

Lines changed: 87 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -823,6 +823,88 @@ void describe('Bedrock converse adapter', () => {
823823
});
824824
});
825825

826+
void it('handles tool use with empty input when streaming', async () => {
827+
const toolOutput: ToolResultContentBlock = {
828+
text: 'additionalToolOutput',
829+
};
830+
const toolExecuteMock = mock.fn<
831+
(input: unknown) => Promise<ToolResultContentBlock>
832+
>(() => Promise.resolve(toolOutput));
833+
const tool: ExecutableTool = {
834+
name: 'toolId',
835+
description: 'tool description',
836+
inputSchema: {
837+
json: {},
838+
},
839+
execute: toolExecuteMock,
840+
};
841+
842+
const event: ConversationTurnEvent = {
843+
...commonEvent,
844+
};
845+
846+
const bedrockClient = new BedrockRuntimeClient();
847+
const bedrockResponseQueue: Array<
848+
ConverseCommandOutput | ConverseStreamCommandOutput
849+
> = [];
850+
const toolUse1 = {
851+
toolUseId: randomUUID().toString(),
852+
name: tool.name,
853+
input: undefined,
854+
};
855+
const toolUse2 = {
856+
toolUseId: randomUUID().toString(),
857+
name: tool.name,
858+
input: '',
859+
};
860+
const toolUseBedrockResponse = mockBedrockResponse(
861+
[
862+
{
863+
toolUse: toolUse1,
864+
},
865+
{
866+
toolUse: toolUse2,
867+
},
868+
],
869+
true
870+
);
871+
bedrockResponseQueue.push(toolUseBedrockResponse);
872+
const content = [
873+
{
874+
text: 'finalResponse',
875+
},
876+
];
877+
const finalBedrockResponse = mockBedrockResponse(content, true);
878+
bedrockResponseQueue.push(finalBedrockResponse);
879+
880+
mock.method(bedrockClient, 'send', () =>
881+
Promise.resolve(bedrockResponseQueue.shift())
882+
);
883+
884+
const adapter = new BedrockConverseAdapter(
885+
event,
886+
[tool],
887+
bedrockClient,
888+
undefined,
889+
messageHistoryRetriever
890+
);
891+
892+
const chunks: Array<StreamingResponseChunk> = await askBedrockWithStreaming(
893+
adapter
894+
);
895+
const responseText = chunks.reduce((acc, next) => {
896+
if (next.contentBlockText) {
897+
acc += next.contentBlockText;
898+
}
899+
return acc;
900+
}, '');
901+
assert.strictEqual(responseText, 'finalResponse');
902+
903+
assert.strictEqual(toolExecuteMock.mock.calls.length, 2);
904+
assert.strictEqual(toolExecuteMock.mock.calls[0].arguments[0], undefined);
905+
assert.strictEqual(toolExecuteMock.mock.calls[1].arguments[0], undefined);
906+
});
907+
826908
void it('throws if tool is duplicated', () => {
827909
assert.throws(
828910
() =>
@@ -1007,19 +1089,21 @@ const mockConverseStreamCommandOutput = (
10071089
},
10081090
},
10091091
});
1010-
const input = JSON.stringify(block.toolUse.input);
1092+
const input = block.toolUse.input
1093+
? JSON.stringify(block.toolUse.input)
1094+
: undefined;
10111095
streamItems.push({
10121096
contentBlockDelta: {
10131097
contentBlockIndex: i,
10141098
delta: {
10151099
toolUse: {
10161100
// simulate chunked input
1017-
input: input.substring(0, 1),
1101+
input: input?.substring(0, 1),
10181102
},
10191103
},
10201104
},
10211105
});
1022-
if (input.length > 1) {
1106+
if (input && input.length > 1) {
10231107
streamItems.push({
10241108
contentBlockDelta: {
10251109
contentBlockIndex: i,

packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -232,8 +232,9 @@ export class BedrockConverseAdapter {
232232
if (chunk.contentBlockDelta.delta?.toolUse) {
233233
if (!chunk.contentBlockDelta.delta.toolUse.input) {
234234
toolUseInput = '';
235+
} else {
236+
toolUseInput += chunk.contentBlockDelta.delta.toolUse.input;
235237
}
236-
toolUseInput += chunk.contentBlockDelta.delta.toolUse.input;
237238
} else if (chunk.contentBlockDelta.delta?.text) {
238239
text += chunk.contentBlockDelta.delta.text;
239240
yield {
@@ -249,7 +250,9 @@ export class BedrockConverseAdapter {
249250
}
250251
} else if (chunk.contentBlockStop) {
251252
if (toolUseBlock) {
252-
toolUseBlock.toolUse.input = JSON.parse(toolUseInput);
253+
if (toolUseInput) {
254+
toolUseBlock.toolUse.input = JSON.parse(toolUseInput);
255+
}
253256
accumulatedAssistantMessage.content?.push(toolUseBlock);
254257
if (
255258
toolUseBlock.toolUse.name &&

0 commit comments

Comments
 (0)