Add s3 source for doc, image, video

Author: Unshure

src/strands/models/bedrock.py

@@ -489,9 +489,17 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
             if "format" in document:
                 result["format"] = document["format"]
 
-            # Handle source
+            # Handle source - supports bytes or s3Location
             if "source" in document:
-                result["source"] = {"bytes": document["source"]["bytes"]}
+                source = document["source"]
+                if "s3Location" in source:
+                    s3_loc = source["s3Location"]
+                    formatted_document_s3: dict[str, Any] = {"uri": s3_loc["uri"]}
+                    if "bucketOwner" in s3_loc:
+                        formatted_document_s3["bucketOwner"] = s3_loc["bucketOwner"]
+                    result["source"] = {"s3Location": formatted_document_s3}
+                elif "bytes" in source:
+                    result["source"] = {"bytes": source["bytes"]}
 
             # Handle optional fields
             if "citations" in document and document["citations"] is not None:
@@ -512,10 +520,16 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
         if "image" in content:
             image = content["image"]
             source = image["source"]
-            formatted_source = {}
-            if "bytes" in source:
-                formatted_source = {"bytes": source["bytes"]}
-            result = {"format": image["format"], "source": formatted_source}
+            formatted_image_source: dict[str, Any] = {}
+            if "s3Location" in source:
+                s3_loc = source["s3Location"]
+                formatted_image_s3: dict[str, Any] = {"uri": s3_loc["uri"]}
+                if "bucketOwner" in s3_loc:
+                    formatted_image_s3["bucketOwner"] = s3_loc["bucketOwner"]
+                formatted_image_source = {"s3Location": formatted_image_s3}
+            elif "bytes" in source:
+                formatted_image_source = {"bytes": source["bytes"]}
+            result = {"format": image["format"], "source": formatted_image_source}
             return {"image": result}
 
         # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ReasoningContentBlock.html
@@ -577,10 +591,16 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
         if "video" in content:
             video = content["video"]
             source = video["source"]
-            formatted_source = {}
-            if "bytes" in source:
-                formatted_source = {"bytes": source["bytes"]}
-            result = {"format": video["format"], "source": formatted_source}
+            formatted_video_source: dict[str, Any] = {}
+            if "s3Location" in source:
+                s3_loc = source["s3Location"]
+                formatted_video_s3: dict[str, Any] = {"uri": s3_loc["uri"]}
+                if "bucketOwner" in s3_loc:
+                    formatted_video_s3["bucketOwner"] = s3_loc["bucketOwner"]
+                formatted_video_source = {"s3Location": formatted_video_s3}
+            elif "bytes" in source:
+                formatted_video_source = {"bytes": source["bytes"]}
+            result = {"format": video["format"], "source": formatted_video_source}
             return {"video": result}
 
         # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CitationsContentBlock.html

src/strands/types/media.py

@@ -7,22 +7,42 @@
 
 from typing import Literal
 
-from typing_extensions import TypedDict
+from typing_extensions import Required, TypedDict
 
 from .citations import CitationsConfig
 
 DocumentFormat = Literal["pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md"]
 """Supported document formats."""
 
 
-class DocumentSource(TypedDict):
+class S3Location(TypedDict, total=False):
+    """A storage location in an Amazon S3 bucket.
+
+    Used by Bedrock to reference media files stored in S3 instead of passing raw bytes.
+
+    - Docs: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_S3Location.html
+
+    Attributes:
+        uri: An object URI starting with `s3://`. Required.
+        bucketOwner: If the bucket belongs to another AWS account, specify that account's ID. Optional.
+    """
+
+    uri: Required[str]
+    bucketOwner: str
+
+
+class DocumentSource(TypedDict, total=False):
     """Contains the content of a document.
 
+    Only one of `bytes` or `s3Location` should be specified.
+
     Attributes:
         bytes: The binary content of the document.
+        s3Location: S3 location of the document (Bedrock only).
     """
 
     bytes: bytes
+    s3Location: S3Location
 
 
 class DocumentContent(TypedDict, total=False):
@@ -45,14 +65,18 @@ class DocumentContent(TypedDict, total=False):
 """Supported image formats."""
 
 
-class ImageSource(TypedDict):
+class ImageSource(TypedDict, total=False):
     """Contains the content of an image.
 
+    Only one of `bytes` or `s3Location` should be specified.
+
     Attributes:
         bytes: The binary content of the image.
+        s3Location: S3 location of the image (Bedrock only).
     """
 
     bytes: bytes
+    s3Location: S3Location
 
 
 class ImageContent(TypedDict):
@@ -71,14 +95,18 @@ class ImageContent(TypedDict):
 """Supported video formats."""
 
 
-class VideoSource(TypedDict):
+class VideoSource(TypedDict, total=False):
     """Contains the content of a video.
 
+    Only one of `bytes` or `s3Location` should be specified.
+
     Attributes:
         bytes: The binary content of the video.
+        s3Location: S3 location of the video (Bedrock only).
     """
 
     bytes: bytes
+    s3Location: S3Location
 
 
 class VideoContent(TypedDict):

tests/strands/agent/test_agent_result.py

@@ -370,4 +370,3 @@ def test__str__empty_interrupts_returns_agent_message(mock_metrics, simple_messa
 
     # Empty list is falsy, should fall through to text content
     assert message_string == "Hello world!\n"
-

tests/strands/models/test_bedrock.py

@@ -1787,8 +1787,8 @@ def test_format_request_filters_image_content_blocks(model, model_id):
     assert "metadata" not in image_block
 
 
-def test_format_request_filters_nested_image_s3_fields(model, model_id):
-    """Test that s3Location is filtered out and only bytes source is preserved."""
+def test_format_request_image_s3_location_only(model, model_id):
+    """Test that image with only s3Location is properly formatted."""
     messages = [
         {
             "role": "user",
@@ -1797,10 +1797,41 @@ def test_format_request_filters_nested_image_s3_fields(model, model_id):
                     "image": {
                         "format": "png",
                         "source": {
-                            "bytes": b"image_data",
-                            "s3Location": {"bucket": "my-bucket", "key": "image.png", "extraField": "filtered"},
+                            "s3Location": {"uri": "s3://my-bucket/image.png"},
                         },
                     }
+                },
+                {
+                    "image": {
+                        "format": "png",
+                        "source": {
+                            "s3Location": {"uri": "s3://my-bucket/image.png", "bucketOwner": "12345"},
+                        },
+                    }
+                },
+            ],
+        }
+    ]
+
+    formatted_request = model._format_request(messages)
+    image_source = formatted_request["messages"][0]["content"][0]["image"]["source"]
+    image_source_with_bucket_owner = formatted_request["messages"][0]["content"][1]["image"]["source"]
+
+    assert image_source == {"s3Location": {"uri": "s3://my-bucket/image.png"}}
+    assert image_source_with_bucket_owner == {"s3Location": {"uri": "s3://my-bucket/image.png", "bucketOwner": "12345"}}
+
+
+def test_format_request_image_bytes_only(model, model_id):
+    """Test that image with only bytes source is properly formatted."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "image": {
+                        "format": "png",
+                        "source": {"bytes": b"image_data"},
+                    }
                 }
             ],
         }
@@ -1810,7 +1841,79 @@ def test_format_request_filters_nested_image_s3_fields(model, model_id):
     image_source = formatted_request["messages"][0]["content"][0]["image"]["source"]
 
     assert image_source == {"bytes": b"image_data"}
-    assert "s3Location" not in image_source
+
+
+def test_format_request_document_s3_location(model, model_id):
+    """Test that document with s3Location is properly formatted."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "document": {
+                        "name": "report.pdf",
+                        "format": "pdf",
+                        "source": {
+                            "s3Location": {"uri": "s3://my-bucket/report.pdf"},
+                        },
+                    }
+                },
+                {
+                    "document": {
+                        "name": "report.pdf",
+                        "format": "pdf",
+                        "source": {
+                            "s3Location": {"uri": "s3://my-bucket/report.pdf", "bucketOwner": "123456789012"},
+                        },
+                    }
+                },
+            ],
+        }
+    ]
+
+    formatted_request = model._format_request(messages)
+    document = formatted_request["messages"][0]["content"][0]["document"]
+    document_with_bucket_owner = formatted_request["messages"][0]["content"][1]["document"]
+
+    assert document["source"] == {"s3Location": {"uri": "s3://my-bucket/report.pdf"}}
+
+    assert document_with_bucket_owner["source"] == {
+        "s3Location": {"uri": "s3://my-bucket/report.pdf", "bucketOwner": "123456789012"}
+    }
+
+
+def test_format_request_video_s3_location(model, model_id):
+    """Test that video with s3Location is properly formatted."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "video": {
+                        "format": "mp4",
+                        "source": {
+                            "s3Location": {"uri": "s3://my-bucket/video.mp4"},
+                        },
+                    }
+                },
+                {
+                    "video": {
+                        "format": "mp4",
+                        "source": {
+                            "s3Location": {"uri": "s3://my-bucket/video.mp4", "bucketOwner": "12345"},
+                        },
+                    }
+                },
+            ],
+        }
+    ]
+
+    formatted_request = model._format_request(messages)
+    video_source = formatted_request["messages"][0]["content"][0]["video"]["source"]
+    video_source_with_bucket_owner = formatted_request["messages"][0]["content"][1]["video"]["source"]
+
+    assert video_source == {"s3Location": {"uri": "s3://my-bucket/video.mp4"}}
+    assert video_source_with_bucket_owner == {"s3Location": {"uri": "s3://my-bucket/video.mp4", "bucketOwner": "12345"}}
 
 
 def test_format_request_filters_document_content_blocks(model, model_id):

tests/strands/tools/mcp/test_mcp_client.py

@@ -632,7 +632,7 @@ def test_call_tool_sync_embedded_nested_base64_textual_mime(mock_transport, mock
 def test_call_tool_sync_embedded_image_blob(mock_transport, mock_session):
     """EmbeddedResource.resource (blob with image MIME) should map to image content."""
     # Read yellow.png file
-    with open("tests_integ/yellow.png", "rb") as image_file:
+    with open("tests_integ/resources/yellow.png", "rb") as image_file:
         png_data = image_file.read()
     payload = base64.b64encode(png_data).decode()
 

tests/strands/types/test_media.py

@@ -0,0 +1,99 @@
+"""Tests for media type definitions."""
+
+from strands.types.media import (
+    DocumentSource,
+    ImageSource,
+    S3Location,
+    VideoSource,
+)
+
+
+class TestS3Location:
+    """Tests for S3Location TypedDict."""
+
+    def test_s3_location_with_uri_only(self):
+        """Test S3Location with only uri field."""
+        s3_loc: S3Location = {"uri": "s3://my-bucket/path/to/file.pdf"}
+
+        assert s3_loc["uri"] == "s3://my-bucket/path/to/file.pdf"
+        assert "bucketOwner" not in s3_loc
+
+    def test_s3_location_with_bucket_owner(self):
+        """Test S3Location with both uri and bucketOwner fields."""
+        s3_loc: S3Location = {
+            "uri": "s3://my-bucket/path/to/file.pdf",
+            "bucketOwner": "123456789012",
+        }
+
+        assert s3_loc["uri"] == "s3://my-bucket/path/to/file.pdf"
+        assert s3_loc["bucketOwner"] == "123456789012"
+
+
+class TestDocumentSource:
+    """Tests for DocumentSource TypedDict."""
+
+    def test_document_source_with_bytes(self):
+        """Test DocumentSource with bytes content."""
+        doc_source: DocumentSource = {"bytes": b"document content"}
+
+        assert doc_source["bytes"] == b"document content"
+        assert "s3Location" not in doc_source
+
+    def test_document_source_with_s3_location(self):
+        """Test DocumentSource with s3Location."""
+        doc_source: DocumentSource = {
+            "s3Location": {
+                "uri": "s3://my-bucket/docs/report.pdf",
+                "bucketOwner": "123456789012",
+            }
+        }
+
+        assert "bytes" not in doc_source
+        assert doc_source["s3Location"]["uri"] == "s3://my-bucket/docs/report.pdf"
+        assert doc_source["s3Location"]["bucketOwner"] == "123456789012"
+
+
+class TestImageSource:
+    """Tests for ImageSource TypedDict."""
+
+    def test_image_source_with_bytes(self):
+        """Test ImageSource with bytes content."""
+        img_source: ImageSource = {"bytes": b"image content"}
+
+        assert img_source["bytes"] == b"image content"
+        assert "s3Location" not in img_source
+
+    def test_image_source_with_s3_location(self):
+        """Test ImageSource with s3Location."""
+        img_source: ImageSource = {
+            "s3Location": {
+                "uri": "s3://my-bucket/images/photo.png",
+            }
+        }
+
+        assert "bytes" not in img_source
+        assert img_source["s3Location"]["uri"] == "s3://my-bucket/images/photo.png"
+
+
+class TestVideoSource:
+    """Tests for VideoSource TypedDict."""
+
+    def test_video_source_with_bytes(self):
+        """Test VideoSource with bytes content."""
+        vid_source: VideoSource = {"bytes": b"video content"}
+
+        assert vid_source["bytes"] == b"video content"
+        assert "s3Location" not in vid_source
+
+    def test_video_source_with_s3_location(self):
+        """Test VideoSource with s3Location."""
+        vid_source: VideoSource = {
+            "s3Location": {
+                "uri": "s3://my-bucket/videos/clip.mp4",
+                "bucketOwner": "987654321098",
+            }
+        }
+
+        assert "bytes" not in vid_source
+        assert vid_source["s3Location"]["uri"] == "s3://my-bucket/videos/clip.mp4"
+        assert vid_source["s3Location"]["bucketOwner"] == "987654321098"