Bedrock - support images in tool results (#2783)

Author: anakin87

integrations/amazon_bedrock/pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.22.0", "boto3>=1.28.57", "aioboto3>=14.0.0"]
+dependencies = ["haystack-ai>=2.23.0", "boto3>=1.28.57", "aioboto3>=14.0.0"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_bedrock#readme"

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

@@ -60,6 +60,24 @@ def _format_tools(tools: list[Tool] | None = None) -> dict[str, Any] | None:
     return {"tools": tool_specs} if tool_specs else None
 
 
+def _convert_image_content_to_bedrock_format(image_content: ImageContent) -> dict[str, Any]:
+    """
+    Convert a Haystack ImageContent to Bedrock format.
+    """
+
+    image_format = image_content.mime_type.split("/")[-1] if image_content.mime_type else None
+    if image_format not in IMAGE_SUPPORTED_FORMATS:
+        err_msg = (
+            f"Unsupported image format: {image_format}. "
+            f"Bedrock supports the following image formats: {IMAGE_SUPPORTED_FORMATS}"
+        )
+        raise ValueError(err_msg)
+
+    source = {"bytes": base64.b64decode(image_content.base64_image)}
+
+    return {"image": {"format": image_format, "source": source}}
+
+
 def _format_tool_call_message(tool_call_message: ChatMessage) -> dict[str, Any]:
     """
     Format a Haystack ChatMessage containing tool calls into Bedrock format.
@@ -94,19 +112,30 @@ def _format_tool_result_message(tool_call_result_message: ChatMessage) -> dict[s
     """
     # Assuming tool call result messages will only contain tool results
     tool_results = []
-    for result in tool_call_result_message.tool_call_results:
-        try:
-            json_result = json.loads(result.result)
-            content = [{"json": json_result}]
-        except json.JSONDecodeError:
-            content = [{"text": result.result}]
+    for tool_call_result in tool_call_result_message.tool_call_results:
+        if isinstance(tool_call_result.result, str):
+            try:
+                json_result = json.loads(tool_call_result.result)
+                content = [{"json": json_result}]
+            except json.JSONDecodeError:
+                content = [{"text": tool_call_result.result}]
+        elif isinstance(tool_call_result.result, list):
+            content = []
+            for item in tool_call_result.result:
+                if isinstance(item, TextContent):
+                    content.append({"text": item.text})
+                elif isinstance(item, ImageContent):
+                    content.append(_convert_image_content_to_bedrock_format(item))
+        else:
+            err_msg = "Unsupported content type in tool call result"
+            raise ValueError(err_msg)
 
         tool_results.append(
             {
                 "toolResult": {
-                    "toolUseId": result.origin.id,
+                    "toolUseId": tool_call_result.origin.id,
                     "content": content,
-                    **({"status": "error"} if result.error else {}),
+                    **({"status": "error"} if tool_call_result.error else {}),
                 }
             }
         )
@@ -217,16 +246,7 @@ def _format_text_image_message(message: ChatMessage) -> dict[str, Any]:
             if message.is_from(ChatRole.ASSISTANT):
                 err_msg = "Image content is not supported for assistant messages"
                 raise ValueError(err_msg)
-
-            image_format = part.mime_type.split("/")[-1] if part.mime_type else None
-            if image_format not in IMAGE_SUPPORTED_FORMATS:
-                err_msg = (
-                    f"Unsupported image format: {image_format}. "
-                    f"Bedrock supports the following image formats: {IMAGE_SUPPORTED_FORMATS}"
-                )
-                raise ValueError(err_msg)
-            source = {"bytes": base64.b64decode(part.base64_image)}
-            bedrock_content_blocks.append({"image": {"format": image_format, "source": source}})
+            bedrock_content_blocks.append(_convert_image_content_to_bedrock_format(part))
 
     return {"role": message.role.value, "content": bedrock_content_blocks}
 

integrations/amazon_bedrock/tests/test_chat_generator.py

@@ -3,10 +3,11 @@
 
 import pytest
 from haystack import Pipeline
+from haystack.components.agents import Agent
 from haystack.components.generators.utils import print_streaming_chunk
 from haystack.components.tools import ToolInvoker
-from haystack.dataclasses import ChatMessage, ChatRole, ImageContent, StreamingChunk, ToolCall
-from haystack.tools import Tool, Toolset
+from haystack.dataclasses import ChatMessage, ChatRole, ImageContent, StreamingChunk, TextContent, ToolCall
+from haystack.tools import Tool, Toolset, create_tool_from_function
 
 from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator
 
@@ -514,6 +515,30 @@ def test_run_with_image_input(self, model_name, test_files_path):
         assert first_reply.text
         assert "apple" in first_reply.text.lower()
 
+    @pytest.mark.parametrize("model_name", MODELS_TO_TEST_WITH_IMAGE_INPUT)
+    def test_live_run_agent_with_images_in_tool_result(self, model_name, test_files_path):
+        def retrieve_image():
+            return [
+                TextContent("Here is the retrieved image."),
+                ImageContent.from_file_path(test_files_path / "apple.jpg", size=(100, 100)),
+            ]
+
+        image_retriever_tool = create_tool_from_function(
+            name="retrieve_image", description="Tool to retrieve an image", function=retrieve_image
+        )
+        image_retriever_tool.outputs_to_string = {"raw_result": True}
+
+        agent = Agent(
+            chat_generator=AmazonBedrockChatGenerator(model=model_name),
+            system_prompt="You are an Agent that can retrieve images and describe them.",
+            tools=[image_retriever_tool],
+        )
+
+        user_message = ChatMessage.from_user("Retrieve the image and describe it in max 5 words.")
+        result = agent.run(messages=[user_message])
+
+        assert "apple" in result["last_message"].text.lower()
+
     @pytest.mark.parametrize("model_name", MODELS_TO_TEST)
     def test_default_inference_with_streaming(self, model_name, chat_messages):
         streaming_callback_called = False

integrations/amazon_bedrock/tests/test_chat_generator_utils.py

@@ -9,6 +9,7 @@
     ImageContent,
     ReasoningContent,
     StreamingChunk,
+    TextContent,
     ToolCall,
     ToolCallDelta,
 )
@@ -120,6 +121,52 @@ def test_format_messages(self):
             {"role": "assistant", "content": [{"text": "The weather in Paris is sunny and 25°C."}]},
         ]
 
+    def test_format_messages_tool_result_with_image(self):
+        base64_image = (
+            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+        )
+
+        messages = [
+            ChatMessage.from_user("Retrieve the image and describe it in max 5 words."),
+            ChatMessage.from_assistant(
+                tool_calls=[ToolCall(id="123", tool_name="image_retriever", arguments={"query": "random query"})]
+            ),
+            ChatMessage.from_tool(
+                tool_result=[
+                    TextContent("Here's the retrieved image"),
+                    ImageContent(base64_image=base64_image, mime_type="image/png"),
+                ],
+                origin=ToolCall(id="123", tool_name="image_retriever", arguments={"query": "random query"}),
+            ),
+            ChatMessage.from_assistant("Beautiful landscape with mountains"),
+        ]
+        formatted_system_prompts, formatted_messages = _format_messages(messages)
+        assert formatted_system_prompts == []
+        assert formatted_messages == [
+            {"role": "user", "content": [{"text": "Retrieve the image and describe it in max 5 words."}]},
+            {
+                "role": "assistant",
+                "content": [
+                    {"toolUse": {"toolUseId": "123", "name": "image_retriever", "input": {"query": "random query"}}}
+                ],
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "toolResult": {
+                            "toolUseId": "123",
+                            "content": [
+                                {"text": "Here's the retrieved image"},
+                                {"image": {"format": "png", "source": {"bytes": base64.b64decode(base64_image)}}},
+                            ],
+                        }
+                    }
+                ],
+            },
+            {"role": "assistant", "content": [{"text": "Beautiful landscape with mountains"}]},
+        ]
+
     def test_format_message_thinking(self):
         assistant_message = ChatMessage.from_assistant(
             "This is a test message.",