Skip to content

Commit 40abb58

Browse files
authored
Extra ChatBedrock unit tests for tool block ordering (#509)
Adding 2 tests for `_format_anthropic_messages()` to verify that `tool_result` and `tool_use` content blocks are reordered correctly for the respective message type.
1 parent 9b2c03e commit 40abb58

File tree

1 file changed

+61
-0
lines changed

1 file changed

+61
-0
lines changed

libs/aws/tests/unit_tests/chat_models/test_bedrock.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -806,6 +806,67 @@ def test__format_anthropic_messages_after_tool_use_no_thinking() -> None:
806806
)
807807

808808

809+
def test__format_anthropic_messages_tool_result_ordering() -> None:
810+
"""Test that tool type content blocks in UserMessage are moved to the beginning."""
811+
human_message = HumanMessage( # type: ignore[misc]
812+
[
813+
{"type": "text", "text": "I have a question about the data."},
814+
{
815+
"type": "tool_result",
816+
"content": "Data analysis result",
817+
"tool_use_id": "tool1"
818+
},
819+
{"type": "text", "text": "Can you explain this result?"},
820+
]
821+
)
822+
823+
system = SystemMessage("You are a helpful assistant") # type: ignore[misc]
824+
ai = AIMessage("Let me think about that") # type: ignore[misc]
825+
826+
messages = [system, human_message, ai]
827+
828+
_, formatted_messages = _format_anthropic_messages(messages)
829+
830+
user_content = formatted_messages[0]["content"]
831+
832+
assert user_content[0]["type"] == "tool_result"
833+
assert user_content[0]["content"] == "Data analysis result"
834+
assert user_content[1]["type"] == "text"
835+
assert user_content[1]["text"] == "I have a question about the data."
836+
assert user_content[2]["type"] == "text"
837+
assert user_content[2]["text"] == "Can you explain this result?"
838+
839+
840+
def test__format_anthropic_messages_tool_use_ordering() -> None:
841+
"""Test that tool type content blocks in AssistantMessage are always moved to the end."""
842+
ai_message = AIMessage( # type: ignore[misc]
843+
[
844+
{"type": "text", "text": "Let me analyze this for you."},
845+
{
846+
"type": "tool_use",
847+
"name": "data_analyzer",
848+
"id": "tool1",
849+
"input": {"data": "sample_data"}
850+
},
851+
{"type": "text", "text": "This will help us understand the pattern."},
852+
]
853+
)
854+
855+
system = SystemMessage("You are a helpful assistant") # type: ignore[misc]
856+
human = HumanMessage("Can you analyze this data?") # type: ignore[misc]
857+
858+
messages = [system, human, ai_message]
859+
860+
_, formatted_messages = _format_anthropic_messages(messages)
861+
862+
assistant_content = formatted_messages[1]["content"]
863+
864+
assert assistant_content[0]["type"] == "text"
865+
assert assistant_content[0]["text"] == "Let me analyze this for you."
866+
assert assistant_content[1]["type"] == "text"
867+
assert assistant_content[1]["text"] == "This will help us understand the pattern."
868+
assert assistant_content[2]["type"] == "tool_use"
869+
assert assistant_content[2]["name"] == "data_analyzer"
809870

810871

811872
def test__format_anthropic_messages_preserves_content_order() -> None:

0 commit comments

Comments
 (0)