diff --git a/src/strands/agent/agent_result.py b/src/strands/agent/agent_result.py index 8f9241a67..b406944da 100644 --- a/src/strands/agent/agent_result.py +++ b/src/strands/agent/agent_result.py @@ -38,29 +38,35 @@ class AgentResult: def __str__(self) -> str: """Get the agent's last message as a string. - This method extracts and concatenates all text content from the final message, ignoring any non-text content - like images or structured data. If there's no text content but structured output is present, it serializes - the structured output instead. + This method extracts and concatenates all text content from the final message, + including text from both "text" blocks and "citationsContent" blocks. + + When structured output exists, its JSON representation is appended to the text. Returns: - The agent's last message as a string. + The agent's last message as a string, including any structured output. """ content_array = self.message.get("content", []) - result = "" + # Collect all text content without modification + text_parts = [] for item in content_array: if isinstance(item, dict): if "text" in item: - result += item.get("text", "") + "\n" + text_parts.append(item.get("text", "")) elif "citationsContent" in item: citations_block = item["citationsContent"] if "content" in citations_block: for content in citations_block["content"]: if isinstance(content, dict) and "text" in content: - result += content.get("text", "") + "\n" + text_parts.append(content.get("text", "")) + + # Join text parts with newline, preserving original content + result = "\n".join(text_parts) + "\n" if text_parts else "" - if not result and self.structured_output: - result = self.structured_output.model_dump_json() + # Append structured output JSON when present + if self.structured_output: + result += self.structured_output.model_dump_json() return result diff --git a/tests/strands/agent/test_agent_result.py b/tests/strands/agent/test_agent_result.py index 6e4c2c91a..364acf13c 100644 --- a/tests/strands/agent/test_agent_result.py +++ b/tests/strands/agent/test_agent_result.py @@ -38,6 +38,46 @@ def empty_message(): return {"role": "assistant", "content": []} +@pytest.fixture +def citations_message(): + """Message with citationsContent block.""" + return { + "role": "assistant", + "content": [ + { + "citationsContent": { + "citations": [ + { + "title": "Source Document", + "location": {"document": {"pageNumber": 1}}, + "sourceContent": [{"text": "source text"}], + } + ], + "content": [{"text": "This is cited text from the document."}], + } + } + ], + } + + +@pytest.fixture +def mixed_text_and_citations_message(): + """Message with both plain text and citationsContent blocks.""" + return { + "role": "assistant", + "content": [ + {"text": "Introduction paragraph"}, + { + "citationsContent": { + "citations": [{"title": "Doc", "location": {}, "sourceContent": []}], + "content": [{"text": "Cited content here."}], + } + }, + {"text": "Conclusion paragraph"}, + ], + } + + def test__init__(mock_metrics, simple_message: Message): """Test that AgentResult can be properly initialized with all required fields.""" stop_reason: StopReason = "end_turn" @@ -99,6 +139,24 @@ def test__str__non_dict_content(mock_metrics): assert message_string == "Valid text\nMore valid text\n" +def test__str__with_citations_content(mock_metrics, citations_message: Message): + """Test that str() extracts text from citationsContent blocks.""" + result = AgentResult(stop_reason="end_turn", message=citations_message, metrics=mock_metrics, state={}) + + message_string = str(result) + assert message_string == "This is cited text from the document.\n" + + +def test__str__mixed_text_and_citations_content(mock_metrics, mixed_text_and_citations_message: Message): + """Test that str() works with both plain text and citationsContent blocks.""" + result = AgentResult( + stop_reason="end_turn", message=mixed_text_and_citations_message, metrics=mock_metrics, state={} + ) + + message_string = str(result) + assert message_string == "Introduction paragraph\nCited content here.\nConclusion paragraph\n" + + def test_to_dict(mock_metrics, simple_message: Message): """Test that to_dict serializes AgentResult correctly.""" result = AgentResult(stop_reason="end_turn", message=simple_message, metrics=mock_metrics, state={"key": "value"}) @@ -185,7 +243,7 @@ def test__init__structured_output_defaults_to_none(mock_metrics, simple_message: def test__str__with_structured_output(mock_metrics, simple_message: Message): - """Test that str() is not affected by structured_output.""" + """Test that str() concatenates text and structured_output.""" structured_output = StructuredOutputModel(name="test", value=42) result = AgentResult( @@ -196,11 +254,9 @@ def test__str__with_structured_output(mock_metrics, simple_message: Message): structured_output=structured_output, ) - # The string representation should only include the message text, not structured output + # str() should concatenate text and structured output message_string = str(result) - assert message_string == "Hello world!\n" - assert "test" not in message_string - assert "42" not in message_string + assert message_string == 'Hello world!\n{"name":"test","value":42,"optional_field":null}' def test__str__empty_message_with_structured_output(mock_metrics, empty_message: Message): @@ -227,59 +283,116 @@ def test__str__empty_message_with_structured_output(mock_metrics, empty_message: assert "optional" in message_string -@pytest.fixture -def citations_message(): - """Message with citationsContent block.""" - return { - "role": "assistant", - "content": [ - { - "citationsContent": { - "citations": [ - { - "title": "Source Document", - "location": {"document": {"pageNumber": 1}}, - "sourceContent": [{"text": "source text"}], - } - ], - "content": [{"text": "This is cited text from the document."}], - } - } - ], - } +def test__str__structured_output_only(): + """Test __str__ with only structured output (no text).""" + structured = StructuredOutputModel(name="test", value=42) + result = AgentResult( + stop_reason="end_turn", + message={"role": "assistant", "content": []}, + metrics=EventLoopMetrics(), + state={}, + structured_output=structured, + ) + # Should return just the structured output JSON + assert str(result) == '{"name":"test","value":42,"optional_field":null}' -@pytest.fixture -def mixed_text_and_citations_message(): - """Message with both plain text and citationsContent blocks.""" - return { - "role": "assistant", - "content": [ - {"text": "Introduction paragraph"}, - { - "citationsContent": { - "citations": [{"title": "Doc", "location": {}, "sourceContent": []}], - "content": [{"text": "Cited content here."}], - } - }, - {"text": "Conclusion paragraph"}, - ], - } +def test__str__both_text_and_structured_output(): + """Test __str__ concatenates text and structured output when both exist.""" + structured = StructuredOutputModel(name="test", value=42) + result = AgentResult( + stop_reason="end_turn", + message={"role": "assistant", "content": [{"text": "Here is the analysis"}]}, + metrics=EventLoopMetrics(), + state={}, + structured_output=structured, + ) + output = str(result) + assert output == 'Here is the analysis\n{"name":"test","value":42,"optional_field":null}' -def test__str__with_citations_content(mock_metrics, citations_message: Message): - """Test that str() extracts text from citationsContent blocks.""" - result = AgentResult(stop_reason="end_turn", message=citations_message, metrics=mock_metrics, state={}) +def test__str__multiple_text_blocks_with_structured_output(): + """Test __str__ with multiple text blocks and structured output.""" + structured = StructuredOutputModel(name="multi", value=100) + result = AgentResult( + stop_reason="end_turn", + message={ + "role": "assistant", + "content": [ + {"text": "First paragraph."}, + {"text": "Second paragraph."}, + ], + }, + metrics=EventLoopMetrics(), + state={}, + structured_output=structured, + ) + output = str(result) + assert output == 'First paragraph.\nSecond paragraph.\n{"name":"multi","value":100,"optional_field":null}' + + +def test__str__non_text_content_only(): + """Test __str__ with only non-text content (e.g., toolUse).""" + result = AgentResult( + stop_reason="tool_use", + message={ + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "123", "name": "test_tool", "input": {}}}], + }, + metrics=EventLoopMetrics(), + state={}, + ) + assert str(result) == "" + + +def test__str__mixed_content_with_structured_output(): + """Test __str__ with mixed content (text + toolUse) and structured output.""" + structured = StructuredOutputModel(name="mixed", value=50) + result = AgentResult( + stop_reason="end_turn", + message={ + "role": "assistant", + "content": [ + {"text": "Processing complete."}, + {"toolUse": {"toolUseId": "456", "name": "helper", "input": {}}}, + ], + }, + metrics=EventLoopMetrics(), + state={}, + structured_output=structured, + ) + output = str(result) + assert output == 'Processing complete.\n{"name":"mixed","value":50,"optional_field":null}' + + +def test__str__citations_with_structured_output(mock_metrics, citations_message: Message): + """Test that str() concatenates citationsContent text and structured_output.""" + structured_output = StructuredOutputModel(name="cited", value=77) + + result = AgentResult( + stop_reason="end_turn", + message=citations_message, + metrics=mock_metrics, + state={}, + structured_output=structured_output, + ) message_string = str(result) - assert message_string == "This is cited text from the document.\n" + assert message_string == 'This is cited text from the document.\n{"name":"cited","value":77,"optional_field":null}' -def test__str__mixed_text_and_citations_content(mock_metrics, mixed_text_and_citations_message: Message): - """Test that str() works with both plain text and citationsContent blocks.""" +def test__str__mixed_text_citations_with_structured_output(mock_metrics, mixed_text_and_citations_message: Message): + """Test that str() handles plain text, citations, and structured output together.""" + structured_output = StructuredOutputModel(name="complex", value=999) + result = AgentResult( - stop_reason="end_turn", message=mixed_text_and_citations_message, metrics=mock_metrics, state={} + stop_reason="end_turn", + message=mixed_text_and_citations_message, + metrics=mock_metrics, + state={}, + structured_output=structured_output, ) message_string = str(result) - assert message_string == "Introduction paragraph\nCited content here.\nConclusion paragraph\n" + expected = 'Introduction paragraph\nCited content here.\nConclusion paragraph\n{"name":"complex","value":999,"optional_field":null}' + assert message_string == expected diff --git a/tests/strands/tools/test_decorator_pep563.py b/tests/strands/tools/test_decorator_pep563.py index 07ec8f2ba..44d9a626a 100644 --- a/tests/strands/tools/test_decorator_pep563.py +++ b/tests/strands/tools/test_decorator_pep563.py @@ -10,10 +10,10 @@ from __future__ import annotations -from typing import Any +from typing import Any, Literal import pytest -from typing_extensions import Literal, TypedDict +from typing_extensions import TypedDict from strands import tool