Refactor to make location generic

Unshure · Unshure · commit 76b1b0b724b4 · 2026-01-28T16:03:37.000-05:00
diff --git a/src/strands/models/bedrock.py b/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 S3Location, SourceLocation
+
 from .._exception_notes import add_exception_note
 from ..event_loop import streaming
 from ..tools import convert_pydantic_to_tool_spec
@@ -407,6 +409,8 @@ def _format_bedrock_messages(self, messages: Messages) -> list[dict[str, Any]]:
 
                 # Format content blocks for Bedrock API compatibility
                 formatted_content = self._format_request_message_content(content_block)
+                if formatted_content is None:
+                    continue
 
                 # Wrap text or image content in guardrailContent if this is the last user message
                 if (
@@ -459,7 +463,19 @@ 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 _format_request_message_content(self, content: ContentBlock) -> dict[str, Any]:
+    def _handle_location(self, location: SourceLocation) -> dict[str, Any] | None:
+        """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 None
+
+    def _format_request_message_content(self, content: ContentBlock) -> dict[str, Any] | None:
         """Format a Bedrock content block.
 
         Bedrock strictly validates content blocks and throws exceptions for unknown fields.
@@ -489,17 +505,17 @@ 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}
+                formatted_document_source: dict[str, Any] | None
+                if "location" in source:
+                    formatted_document_source = self._handle_location(source["location"])
+                    if formatted_document_source is None:
+                        return None
                 elif "bytes" in source:
-                    result["source"] = {"bytes": source["bytes"]}
+                    formatted_document_source = {"bytes": source["bytes"]}
+                result["source"] = formatted_document_source
 
             # Handle optional fields
             if "citations" in document and document["citations"] is not None:
@@ -520,13 +536,11 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
         if "image" in content:
             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}
+            formatted_image_source: dict[str, Any] | None
+            if "location" in source:
+                formatted_image_source = self._handle_location(source["location"])
+                if formatted_image_source is None:
+                    return None
             elif "bytes" in source:
                 formatted_image_source = {"bytes": source["bytes"]}
             result = {"format": image["format"], "source": formatted_image_source}
@@ -564,9 +578,12 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
                     # Handle json field since not in ContentBlock but valid in ToolResultContent
                     formatted_content.append({"json": tool_result_content["json"]})
                 else:
-                    formatted_content.append(
-                        self._format_request_message_content(cast(ContentBlock, tool_result_content))
+                    formatted_message_content = self._format_request_message_content(
+                        cast(ContentBlock, tool_result_content)
                     )
+                    if formatted_message_content is None:
+                        continue
+                    formatted_content.append(formatted_message_content)
 
             result = {
                 "content": formatted_content,
@@ -591,13 +608,11 @@ def _format_request_message_content(self, content: ContentBlock) -> dict[str, An
         if "video" in content:
             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}
+            formatted_video_source: dict[str, Any] | None
+            if "location" in source:
+                formatted_video_source = self._handle_location(source["location"])
+                if formatted_video_source is None:
+                    return None
             elif "bytes" in source:
                 formatted_video_source = {"bytes": source["bytes"]}
             result = {"format": video["format"], "source": formatted_video_source}
diff --git a/src/strands/types/media.py b/src/strands/types/media.py
@@ -5,7 +5,7 @@
 - Bedrock docs: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Types_Amazon_Bedrock_Runtime.html
 """
 
-from typing import Literal
+from typing import Literal, TypeAlias
 
 from typing_extensions import Required, TypedDict
 
@@ -15,34 +15,50 @@
 """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.
     """
 
+    # mypy doesn't like overriding this field since its a subclass, but since its just a literal string, this is fine.
+
+    type: Literal["s3"]  # type: ignore[misc]
     uri: Required[str]
     bucketOwner: str
 
 
+SourceLocation: TypeAlias = Location | S3Location
+
+
 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).
+        location: Location of the document.
     """
 
     bytes: bytes
-    s3Location: S3Location
+    location: SourceLocation
 
 
 class DocumentContent(TypedDict, total=False):
@@ -72,11 +88,11 @@ class ImageSource(TypedDict, total=False):
 
     Attributes:
         bytes: The binary content of the image.
-        s3Location: S3 location of the image (Bedrock only).
+        location: Location of the image.
     """
 
     bytes: bytes
-    s3Location: S3Location
+    location: SourceLocation
 
 
 class ImageContent(TypedDict):
@@ -102,11 +118,11 @@ class VideoSource(TypedDict, total=False):
 
     Attributes:
         bytes: The binary content of the video.
-        s3Location: S3 location of the video (Bedrock only).
+        location: Location of the video.
     """
 
     bytes: bytes
-    s3Location: S3Location
+    location: SourceLocation
 
 
 class VideoContent(TypedDict):
diff --git a/tests/strands/agent/hooks/test_events.py b/tests/strands/agent/hooks/test_events.py
@@ -206,8 +206,6 @@ def test_invocation_state_is_available_in_model_call_events(agent):
     assert after_event.invocation_state["request_id"] == "req-456"
 
 
-
-
 def test_before_invocation_event_messages_default_none(agent):
     """Test that BeforeInvocationEvent.messages defaults to None for backward compatibility."""
     event = BeforeInvocationEvent(agent=agent)
diff --git a/tests/strands/models/test_bedrock.py b/tests/strands/models/test_bedrock.py
@@ -1,3 +1,5 @@
+import copy
+import logging
 import os
 import sys
 import traceback
@@ -1519,7 +1521,6 @@ async def test_add_note_on_validation_exception_throughput(bedrock_client, model
 @pytest.mark.asyncio
 async def test_stream_logging(bedrock_client, model, messages, caplog, alist):
     """Test that stream method logs debug messages at the expected stages."""
-    import logging
 
     # Set the logger to debug level to capture debug messages
     caplog.set_level(logging.DEBUG, logger="strands.models.bedrock")
@@ -1797,28 +1798,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 +1845,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 +1854,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",
+                            },
                         },
                     }
                 },
@@ -1882,25 +1877,67 @@ def test_format_request_document_s3_location(model, model_id):
     }
 
 
-def test_format_request_video_s3_location(model, model_id):
-    """Test that video with s3Location is properly formatted."""
+def test_format_request_unsupported_location(model, caplog):
+    """Test that document with s3Location is properly formatted."""
+
+    caplog.set_level(logging.WARNING, logger="strands.models.bedrock")
+
     messages = [
         {
             "role": "user",
             "content": [
+                {"text": "Hello!"},
+                {
+                    "document": {
+                        "name": "report.pdf",
+                        "format": "pdf",
+                        "source": {
+                            "location": {
+                                "type": "other",
+                            },
+                        },
+                    }
+                },
                 {
                     "video": {
                         "format": "mp4",
                         "source": {
-                            "s3Location": {"uri": "s3://my-bucket/video.mp4"},
+                            "location": {
+                                "type": "other",
+                            },
+                        },
+                    }
+                },
+                {
+                    "image": {
+                        "format": "png",
+                        "source": {
+                            "location": {
+                                "type": "other",
+                            },
                         },
                     }
                 },
+            ],
+        }
+    ]
+
+    formatted_request = model._format_request(messages)
+    assert len(formatted_request["messages"][0]["content"]) == 1
+    assert "Non s3 location sources are not supported by Bedrock, skipping content block" in caplog.text
+
+
+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", "bucketOwner": "12345"},
+                            "location": {"type": "s3", "uri": "s3://my-bucket/video.mp4"},
                         },
                     }
                 },
@@ -1910,10 +1947,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):
@@ -2413,7 +2448,6 @@ def test_inject_cache_point_skipped_for_non_claude(bedrock_client):
 
 def test_format_bedrock_messages_does_not_mutate_original(bedrock_client):
     """Test that _format_bedrock_messages does not mutate original messages."""
-    import copy
 
     model = BedrockModel(
         model_id="us.anthropic.claude-sonnet-4-20250514-v1:0", cache_config=CacheConfig(strategy="auto")
diff --git a/tests_integ/test_bedrock_s3_location.py b/tests_integ/test_bedrock_s3_location.py
