Skip to content

Commit 3f1c5c5

Browse files
anakin87anakin87
authored
fix: AmazonBedrockChatGenerator - fix bug with streaming + tool calls with no arguments (#2121)
* fix: AmazonBedrockChatGenerator - fix bug with streaming + tool calls with no arguments * uncomment lines
1 parent a679e97 commit 3f1c5c5

3 files changed

Lines changed: 310 additions & 2 deletions

File tree

‎integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,7 @@ def _convert_streaming_chunks_to_chat_message(chunks: List[StreamingChunk]) -> C
447447
# Convert accumulated tool call data into ToolCall objects
448448
for call_data in tool_call_data.values():
449449
try:
450-
arguments = json.loads(call_data["arguments"])
450+
arguments = json.loads(call_data.get("arguments", "{}")) if call_data.get("arguments") else {}
451451
tool_calls.append(ToolCall(id=call_data["id"], tool_name=call_data["name"], arguments=arguments))
452452
except json.JSONDecodeError:
453453
logger.warning(

‎integrations/amazon_bedrock/tests/test_chat_generator.py‎

Lines changed: 55 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from haystack import Pipeline
55
from haystack.components.generators.utils import print_streaming_chunk
66
from haystack.components.tools import ToolInvoker
7-
from haystack.dataclasses import ChatMessage, ChatRole, ImageContent, StreamingChunk
7+
from haystack.dataclasses import ChatMessage, ChatRole, ImageContent, StreamingChunk, ToolCall
88
from haystack.tools import Tool
99

1010
from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator
@@ -29,6 +29,21 @@
2929
]
3030

3131

32+
def hello_world():
33+
return "Hello, World!"
34+
35+
36+
@pytest.fixture
37+
def tool_with_no_parameters():
38+
tool = Tool(
39+
name="hello_world",
40+
description="This prints hello world",
41+
parameters={"properties": {}, "type": "object"},
42+
function=hello_world,
43+
)
44+
return tool
45+
46+
3247
def weather(city: str):
3348
"""Get weather for a given city."""
3449
return f"The weather in {city} is sunny and 32°C"
@@ -430,6 +445,45 @@ def test_live_run_with_multi_tool_calls_streaming(self, model_name, tools):
430445
assert "paris" in final_message.text.lower()
431446
assert "berlin" in final_message.text.lower()
432447

448+
@pytest.mark.parametrize("model_name", [STREAMING_TOOL_MODELS[1]]) # just one model is enough
449+
def test_live_run_with_tool_with_no_args_streaming(self, tool_with_no_parameters, model_name):
450+
"""
451+
Integration test that the AmazonBedrockChatGenerator component can run with a tool that has no arguments and
452+
streaming.
453+
"""
454+
initial_messages = [ChatMessage.from_user("Print Hello World using the print hello world tool.")]
455+
component = AmazonBedrockChatGenerator(
456+
model=model_name, tools=[tool_with_no_parameters], streaming_callback=print_streaming_chunk
457+
)
458+
results = component.run(messages=initial_messages)
459+
460+
assert len(results["replies"]) == 1
461+
message = results["replies"][0]
462+
463+
# # this is Claude thinking message prior to tool call
464+
# assert message.text is not None
465+
466+
# now we have the tool call
467+
assert message.tool_calls
468+
tool_call = message.tool_call
469+
assert isinstance(tool_call, ToolCall)
470+
assert tool_call.id is not None
471+
assert tool_call.tool_name == "hello_world"
472+
assert tool_call.arguments == {}
473+
assert message.meta["finish_reason"] == "tool_use"
474+
475+
new_messages = [
476+
*initial_messages,
477+
message,
478+
ChatMessage.from_tool(tool_result="Hello World!", origin=tool_call),
479+
]
480+
results = component.run(new_messages)
481+
assert len(results["replies"]) == 1
482+
final_message = results["replies"][0]
483+
assert not final_message.tool_calls
484+
assert len(final_message.text) > 0
485+
assert "hello" in final_message.text.lower()
486+
433487
@pytest.mark.parametrize("model_name", [MODELS_TO_TEST_WITH_TOOLS[0]]) # just one model is enough
434488
def test_pipeline_with_amazon_bedrock_chat_generator(self, model_name, tools):
435489
"""

‎integrations/amazon_bedrock/tests/test_chat_generator_utils.py‎

Lines changed: 254 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from haystack.tools import Tool
66

77
from haystack_integrations.components.generators.amazon_bedrock.chat.utils import (
8+
_convert_streaming_chunks_to_chat_message,
89
_format_messages,
910
_format_text_image_message,
1011
_format_tools,
@@ -555,3 +556,256 @@ def test_callback(chunk: StreamingChunk):
555556
),
556557
]
557558
assert replies == expected_messages
559+
560+
def test_convert_streaming_chunks_to_chat_message_tool_call_with_empty_arguments(self):
561+
chunks = [
562+
StreamingChunk(
563+
content="Certainly! I",
564+
meta={
565+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
566+
"index": 0,
567+
"tool_calls": None,
568+
"finish_reason": None,
569+
"received_at": "2025-07-31T08:46:07.072764+00:00",
570+
},
571+
),
572+
StreamingChunk(
573+
content=" can help",
574+
meta={
575+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
576+
"index": 0,
577+
"tool_calls": None,
578+
"finish_reason": None,
579+
"received_at": "2025-07-31T08:46:07.111264+00:00",
580+
},
581+
),
582+
StreamingChunk(
583+
content=" you print",
584+
meta={
585+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
586+
"index": 0,
587+
"tool_calls": None,
588+
"finish_reason": None,
589+
"received_at": "2025-07-31T08:46:07.162575+00:00",
590+
},
591+
),
592+
StreamingChunk(
593+
content=' "Hello World"',
594+
meta={
595+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
596+
"index": 0,
597+
"tool_calls": None,
598+
"finish_reason": None,
599+
"received_at": "2025-07-31T08:46:07.215535+00:00",
600+
},
601+
),
602+
StreamingChunk(
603+
content=" using the available",
604+
meta={
605+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
606+
"index": 0,
607+
"tool_calls": None,
608+
"finish_reason": None,
609+
"received_at": "2025-07-31T08:46:07.270642+00:00",
610+
},
611+
),
612+
StreamingChunk(
613+
content=' "',
614+
meta={
615+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
616+
"index": 0,
617+
"tool_calls": None,
618+
"finish_reason": None,
619+
"received_at": "2025-07-31T08:46:07.349415+00:00",
620+
},
621+
),
622+
StreamingChunk(
623+
content='hello_world" tool',
624+
meta={
625+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
626+
"index": 0,
627+
"tool_calls": None,
628+
"finish_reason": None,
629+
"received_at": "2025-07-31T08:46:07.426891+00:00",
630+
},
631+
),
632+
StreamingChunk(
633+
content=". This tool is",
634+
meta={
635+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
636+
"index": 0,
637+
"tool_calls": None,
638+
"finish_reason": None,
639+
"received_at": "2025-07-31T08:46:07.495910+00:00",
640+
},
641+
),
642+
StreamingChunk(
643+
content=' designed to print "',
644+
meta={
645+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
646+
"index": 0,
647+
"tool_calls": None,
648+
"finish_reason": None,
649+
"received_at": "2025-07-31T08:46:07.527426+00:00",
650+
},
651+
),
652+
StreamingChunk(
653+
content='Hello World" an',
654+
meta={
655+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
656+
"index": 0,
657+
"tool_calls": None,
658+
"finish_reason": None,
659+
"received_at": "2025-07-31T08:46:07.590629+00:00",
660+
},
661+
),
662+
StreamingChunk(
663+
content="d doesn't require any parameters",
664+
meta={
665+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
666+
"index": 0,
667+
"tool_calls": None,
668+
"finish_reason": None,
669+
"received_at": "2025-07-31T08:46:07.682261+00:00",
670+
},
671+
),
672+
StreamingChunk(
673+
content=". Let",
674+
meta={
675+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
676+
"index": 0,
677+
"tool_calls": None,
678+
"finish_reason": None,
679+
"received_at": "2025-07-31T08:46:07.790526+00:00",
680+
},
681+
),
682+
StreamingChunk(
683+
content="'s go ahead an",
684+
meta={
685+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
686+
"index": 0,
687+
"tool_calls": None,
688+
"finish_reason": None,
689+
"received_at": "2025-07-31T08:46:07.845332+00:00",
690+
},
691+
),
692+
StreamingChunk(
693+
content="d use",
694+
meta={
695+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
696+
"index": 0,
697+
"tool_calls": None,
698+
"finish_reason": None,
699+
"received_at": "2025-07-31T08:46:07.990588+00:00",
700+
},
701+
),
702+
StreamingChunk(
703+
content=" it.",
704+
meta={
705+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
706+
"index": 0,
707+
"tool_calls": None,
708+
"finish_reason": None,
709+
"received_at": "2025-07-31T08:46:07.994309+00:00",
710+
},
711+
),
712+
StreamingChunk(
713+
content="",
714+
meta={
715+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
716+
"received_at": "2025-07-31T08:46:08.359127+00:00",
717+
},
718+
),
719+
StreamingChunk(
720+
content="",
721+
meta={
722+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
723+
"index": 0,
724+
"tool_calls": [
725+
{
726+
"index": 1,
727+
"id": "tooluse_QZlUqTveTwyUaCQGQbWP6g",
728+
"function": {"arguments": "", "name": "hello_world"},
729+
"type": "function",
730+
}
731+
],
732+
"finish_reason": None,
733+
"received_at": "2025-07-31T08:46:08.359912+00:00",
734+
},
735+
),
736+
StreamingChunk(
737+
content="",
738+
meta={
739+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
740+
"index": 0,
741+
"tool_calls": [
742+
{"index": 1, "id": None, "function": {"arguments": "", "name": None}, "type": "function"}
743+
],
744+
"finish_reason": None,
745+
"received_at": "2025-07-31T08:46:08.361612+00:00",
746+
},
747+
),
748+
StreamingChunk(
749+
content="",
750+
meta={
751+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
752+
"received_at": "2025-07-31T08:46:08.591668+00:00",
753+
},
754+
),
755+
StreamingChunk(
756+
content="",
757+
meta={
758+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
759+
"index": 0,
760+
"tool_calls": None,
761+
"finish_reason": "tool_use",
762+
"received_at": "2025-07-31T08:46:08.592175+00:00",
763+
},
764+
index=None,
765+
tool_calls=None,
766+
tool_call_result=None,
767+
start=False,
768+
finish_reason=None,
769+
),
770+
StreamingChunk(
771+
content="",
772+
meta={
773+
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
774+
"index": 0,
775+
"tool_calls": None,
776+
"finish_reason": None,
777+
"received_at": "2025-07-31T08:46:08.596700+00:00",
778+
"usage": {"prompt_tokens": 349, "completion_tokens": 84, "total_tokens": 433},
779+
},
780+
index=None,
781+
tool_calls=None,
782+
tool_call_result=None,
783+
start=False,
784+
finish_reason=None,
785+
),
786+
]
787+
788+
message = _convert_streaming_chunks_to_chat_message(chunks)
789+
790+
# Verify the message content
791+
assert message.text == (
792+
'Certainly! I can help you print "Hello World" using the available "hello_world" tool. This tool is '
793+
"designed to print \"Hello World\" and doesn't require any parameters. Let's go ahead and use it."
794+
)
795+
796+
# Verify tool calls
797+
assert len(message.tool_calls) == 1
798+
tool_call = message.tool_calls[0]
799+
assert tool_call.id == "tooluse_QZlUqTveTwyUaCQGQbWP6g"
800+
assert tool_call.tool_name == "hello_world"
801+
assert tool_call.arguments == {}
802+
803+
# Verify meta information
804+
assert message._meta["model"] == "anthropic.claude-3-5-sonnet-20240620-v1:0"
805+
assert message._meta["index"] == 0
806+
assert message._meta["finish_reason"] == "tool_use"
807+
assert message._meta["usage"] == {
808+
"completion_tokens": 84,
809+
"prompt_tokens": 349,
810+
"total_tokens": 433,
811+
}

0 commit comments

Comments
 (0)