Refactor to make location generic

Author: Unshure

src/strands/models/bedrock.py

@@ -17,6 +17,8 @@
 from pydantic import BaseModel
 from typing_extensions import TypedDict, Unpack, override
 
+from strands.types.media import Location, S3Location
+
 from .._exception_notes import add_exception_note
 from ..event_loop import streaming
 from ..tools import convert_pydantic_to_tool_spec
@@ -459,6 +461,18 @@ def _should_include_tool_result_status(self) -> bool:
         else:  # "auto"
             return any(model in self.config["model_id"] for model in _MODELS_INCLUDE_STATUS)
 
+    def _handle_location(self, location: Location) -> dict[str, Any]:
+        """Convert location content block to Bedrock format if its an S3Location."""
+        if location["type"] == "s3":
+            s3_location = cast(S3Location, location)
+            formatted_document_s3: dict[str, Any] = {"uri": s3_location["uri"]}
+            if "bucketOwner" in s3_location:
+                formatted_document_s3["bucketOwner"] = s3_location["bucketOwner"]
+            return {"s3Location": formatted_document_s3}
+        else:
+            logger.warning("Non s3 location sources are not supported by Bedrock, skipping content block")
+            return {}
+
     def _format_request_message_content(self, content: ContentBlock) -> dict[str, Any]:
         """Format a Bedrock content block.
 
@@ -489,15 +503,11 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
             if "format" in document:
                 result["format"] = document["format"]
 
-            # Handle source - supports bytes or s3Location
+            # Handle source - supports bytes or location
             if "source" in document:
                 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}
+                if "location" in source:
+                    result["source"] = self._handle_location(source["location"])
                 elif "bytes" in source:
                     result["source"] = {"bytes": source["bytes"]}
 
@@ -521,12 +531,8 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
             image = content["image"]
             source = image["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}
+            if "location" in source:
+                formatted_image_source = self._handle_location(source["location"])
             elif "bytes" in source:
                 formatted_image_source = {"bytes": source["bytes"]}
             result = {"format": image["format"], "source": formatted_image_source}
@@ -592,12 +598,8 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
             video = content["video"]
             source = video["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}
+            if "location" in source:
+                formatted_video_source = self._handle_location(source["location"])
             elif "bytes" in source:
                 formatted_video_source = {"bytes": source["bytes"]}
             result = {"format": video["format"], "source": formatted_video_source}

src/strands/types/media.py

@@ -15,18 +15,29 @@
 """Supported document formats."""
 
 
-class S3Location(TypedDict, total=False):
+class Location(TypedDict, total=False):
+    """A location for a document.
+
+    This type is a generic location for a document. Its usage is determined by the underlying model provider
+    """
+
+    type: Required[str]
+
+
+class S3Location(Location, 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:
+        type: s3
         uri: An object URI starting with `s3://`. Required.
         bucketOwner: If the bucket belongs to another AWS account, specify that account's ID. Optional.
     """
 
+    type: Literal["s3"]
     uri: Required[str]
     bucketOwner: str
 
@@ -38,11 +49,11 @@ class DocumentSource(TypedDict, total=False):
 
     Attributes:
         bytes: The binary content of the document.
-        s3Location: S3 location of the document (Bedrock only).
+        location: Location of the document.
     """
 
     bytes: bytes
-    s3Location: S3Location
+    location: Location | S3Location
 
 
 class DocumentContent(TypedDict, total=False):

tests/strands/models/test_bedrock.py

@@ -1797,28 +1797,18 @@ def test_format_request_image_s3_location_only(model, model_id):
                     "image": {
                         "format": "png",
                         "source": {
-                            "s3Location": {"uri": "s3://my-bucket/image.png"},
+                            "location": {"type": "s3", "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):
@@ -1854,7 +1844,7 @@ def test_format_request_document_s3_location(model, model_id):
                         "name": "report.pdf",
                         "format": "pdf",
                         "source": {
-                            "s3Location": {"uri": "s3://my-bucket/report.pdf"},
+                            "location": {"type": "s3", "uri": "s3://my-bucket/report.pdf"},
                         },
                     }
                 },
@@ -1863,7 +1853,11 @@ def test_format_request_document_s3_location(model, model_id):
                         "name": "report.pdf",
                         "format": "pdf",
                         "source": {
-                            "s3Location": {"uri": "s3://my-bucket/report.pdf", "bucketOwner": "123456789012"},
+                            "location": {
+                                "type": "s3",
+                                "uri": "s3://my-bucket/report.pdf",
+                                "bucketOwner": "123456789012",
+                            },
                         },
                     }
                 },
@@ -1892,15 +1886,7 @@ def test_format_request_video_s3_location(model, model_id):
                     "video": {
                         "format": "mp4",
                         "source": {
-                            "s3Location": {"uri": "s3://my-bucket/video.mp4"},
-                        },
-                    }
-                },
-                {
-                    "video": {
-                        "format": "mp4",
-                        "source": {
-                            "s3Location": {"uri": "s3://my-bucket/video.mp4", "bucketOwner": "12345"},
+                            "location": {"type": "s3", "uri": "s3://my-bucket/video.mp4"},
                         },
                     }
                 },
@@ -1910,10 +1896,8 @@ def test_format_request_video_s3_location(model, model_id):
 
     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_integ/test_bedrock_s3_location.py

@@ -120,14 +120,14 @@ def test_document_s3_location(s3_document, account_id):
                     "document": {
                         "format": "pdf",
                         "name": "letter",
-                        "source": {"s3Location": {"uri": s3_document, "bucketOwner": account_id}},
+                        "source": {"location": {"type": "s3", "uri": s3_document, "bucketOwner": account_id}},
                     },
                 },
             ],
         },
     ]
 
-    agent = Agent(model=BedrockModel(model_id="amazon.nova-2-lite-v1:0", region_name="us-west-2"))
+    agent = Agent(model=BedrockModel(model_id="us.amazon.nova-2-lite-v1:0", region_name="us-west-2"))
     result = agent(messages)
 
     assert "amazon" in str(result).lower()
@@ -143,14 +143,14 @@ def test_image_s3_location(s3_image):
                 {
                     "image": {
                         "format": "png",
-                        "source": {"s3Location": {"uri": s3_image}},
+                        "source": {"location": {"type": "s3", "uri": s3_image}},
                     },
                 },
             ],
         },
     ]
 
-    agent = Agent(model=BedrockModel(model_id="amazon.nova-2-lite-v1:0", region_name="us-west-2"))
+    agent = Agent(model=BedrockModel(model_id="us.amazon.nova-2-lite-v1:0", region_name="us-west-2"))
     result = agent(messages)
 
     assert "yellow" in str(result).lower()
@@ -163,7 +163,7 @@ def test_video_s3_location(s3_video):
             "role": "user",
             "content": [
                 {"text": "Describe the colors is in this video?"},
-                {"video": {"format": "mp4", "source": {"s3Location": {"uri": s3_video}}}},
+                {"video": {"format": "mp4", "source": {"location": {"type": "s3", "uri": s3_video}}}},
             ],
         },
     ]