diff --git a/README.md b/README.md index b8ef9ce7..91555286 100644 --- a/README.md +++ b/README.md @@ -123,7 +123,7 @@ Below is a comprehensive table of all available tools, how to use them with an a | environment | `agent.tool.environment(action="list", prefix="AWS_")` | Managing environment variables, configuration management | | generate_image_stability | `agent.tool.generate_image_stability(prompt="A tranquil pool")` | Creating images using Stability AI models | | generate_image | `agent.tool.generate_image(prompt="A sunset over mountains")` | Creating AI-generated images for various applications | -| image_reader | `agent.tool.image_reader(image_path="path/to/image.jpg")` | Processing and reading image files for AI analysis | +| image_reader | `agent.tool.image_reader(image_path="path/to/image.jpg") # or use URL: "https://image.png"` | Processing and reading image files for AI analysis | | journal | `agent.tool.journal(action="write", content="Today's progress notes")` | Creating structured logs, maintaining documentation | | think | `agent.tool.think(thought="Complex problem to analyze", cycle_count=3)` | Advanced reasoning, multi-step thinking processes | | load_tool | `agent.tool.load_tool(path="path/to/custom_tool.py", name="custom_tool")` | Dynamically loading custom tools and extensions | diff --git a/src/strands_tools/image_reader.py b/src/strands_tools/image_reader.py index 240ae426..dafbfa18 100644 --- a/src/strands_tools/image_reader.py +++ b/src/strands_tools/image_reader.py @@ -33,21 +33,29 @@ # With user directory path result = agent.tool.image_reader(image_path="~/Documents/images/photo.png") + +# With custom url +result = agent.tool.image_reader(image_path="https://image.png") ``` See the image_reader function docstring for more details on parameters and return format. """ import os +from io import BytesIO from os.path import expanduser from typing import Any +from urllib.parse import urlparse +import requests from PIL import Image from strands.types.tools import ToolResult, ToolUse TOOL_SPEC = { "name": "image_reader", - "description": "Reads an image file from a given path and returns it in the format required for the Converse API", + "description": ( + "Reads an image file from a given path or url and returns it in the format required for the Converse API" + ), "inputSchema": { "json": { "type": "object", @@ -114,6 +122,7 @@ def image_reader(tool: ToolUse, **kwargs: Any) -> ToolResult: - If the image format is not recognized, it defaults to PNG - The function validates file existence before attempting to read - User paths with tilde (~) are automatically expanded + - Supports direct url """ try: tool_use_id = tool["toolUseId"] @@ -126,18 +135,34 @@ def image_reader(tool: ToolUse, **kwargs: Any) -> ToolResult: "content": [{"text": "File path is required"}], } - file_path = expanduser(tool_input.get("image_path")) + image_path = tool_input.get("image_path") + parsed_image_path = urlparse(image_path) + is_url = all([parsed_image_path.scheme, parsed_image_path.netloc]) + + file_bytes = BytesIO() + if is_url: + file = BytesIO(requests.get(image_path).content) + file_path = file + file_bytes = file.getvalue() + else: + file_path = expanduser(image_path) + if not os.path.exists(file_path): + return { + "toolUseId": tool_use_id, + "status": "error", + "content": [{"text": f"File not found at path: {file_path}"}], + } + + with open(file_path, "rb") as file: + file_bytes = file.read() - if not os.path.exists(file_path): + if not file_path: return { "toolUseId": tool_use_id, "status": "error", - "content": [{"text": f"File not found at path: {file_path}"}], + "content": [{"text": "File path not found"}], } - with open(file_path, "rb") as file: - file_bytes = file.read() - # Handle image files using PIL with Image.open(file_path) as img: image_format = img.format.lower() @@ -149,6 +174,12 @@ def image_reader(tool: ToolUse, **kwargs: Any) -> ToolResult: "status": "success", "content": [{"image": {"format": image_format, "source": {"bytes": file_bytes}}}], } + except requests.exceptions.RequestException as e: + return { + "toolUseId": tool_use_id, + "status": "error", + "content": [{"text": f"Error fetching file from url: {str(e)}"}], + } except Exception as e: return { "toolUseId": tool_use_id, diff --git a/tests/test_image_reader.py b/tests/test_image_reader.py index ac7d7512..13f7647d 100644 --- a/tests/test_image_reader.py +++ b/tests/test_image_reader.py @@ -17,6 +17,13 @@ def test_image_path(): return os.path.expanduser("~/test_image.jpg") +@pytest.fixture +def test_image_url(): + """Return url to a test image file.""" + # This is a placeholder - we'll use a mock instead of creating real files + return "https://image.png" + + @pytest.fixture def test_video_path(): """Return path to a test video file.""" @@ -64,6 +71,38 @@ def test_image_reader_with_image(mock_pil_open, mock_open, mock_exists, test_ima assert result["content"][0]["image"]["source"]["bytes"] == b"fake_image_bytes" +@patch("strands_tools.image_reader.requests.get") +@patch("strands_tools.image_reader.Image.open") +def test_image_reader_with_url(mock_pil_open, mock_requests_get, test_image_url): + """Test the image_reader tool with an image URL.""" + + # Mock requests.get to simulate downloading the image + mock_response = MagicMock() + mock_response.content = b"fake_url_image_bytes" + mock_requests_get.return_value = mock_response + + # Mock PIL Image + mock_img = MagicMock() + mock_img.__enter__.return_value.format = "PNG" + mock_pil_open.return_value = mock_img + + # Call the tool directly + tool_use = { + "toolUseId": "test-tool-use-id", + "input": {"image_path": test_image_url}, + } + + result = image_reader.image_reader(tool=tool_use) + + # Verify result structure + assert result["toolUseId"] == "test-tool-use-id" + assert result["status"] == "success" + assert "image" in result["content"][0] + image_data = result["content"][0]["image"] + assert image_data["format"].lower() == "png" + assert image_data["source"]["bytes"] == b"fake_url_image_bytes" + + @patch("strands_tools.image_reader.os.path.exists") def test_image_reader_file_not_found(mock_exists): """Test error handling when file does not exist."""