fix: move to complete mapping like other providers

Author: dbschmigelski

src/strands/models/bedrock.py

@@ -17,7 +17,7 @@
 
 from ..event_loop import streaming
 from ..tools import convert_pydantic_to_tool_spec
-from ..types.content import ContentBlock, Message, Messages, _ContentBlockType
+from ..types.content import ContentBlock, Messages
 from ..types.exceptions import (
     ContextWindowOverflowException,
     ModelThrottledException,
@@ -43,83 +43,6 @@
     "anthropic.claude",
 ]
 
-# Allowed fields for each Bedrock content block type to prevent validation exceptions
-# Bedrock strictly validates content blocks and throws exceptions for unknown fields
-# https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlock.html
-_BEDROCK_CONTENT_BLOCK_FIELDS: dict[_ContentBlockType, set[str]] = {
-    "image": {
-        "format",
-        "source",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageBlock.html
-    "toolResult": {
-        "content",
-        "toolUseId",
-        "status",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolResultBlock.html
-    "toolUse": {
-        "input",
-        "name",
-        "toolUseId",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolUseBlock.html
-    "document": {
-        "name",
-        "source",
-        "citations",
-        "context",
-        "format",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_DocumentBlock.html
-    "video": {
-        "format",
-        "source",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_VideoBlock.html
-    "reasoningContent": {
-        "reasoningText",
-        "redactedContent",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ReasoningContentBlock.html
-    "citationsContent": {
-        "citations",
-        "content",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CitationsContentBlock.html
-    "cachePoint": {"type"},  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CachePointBlock.html
-    "guardContent": {
-        "image",
-        "text",
-    },  # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_GuardrailConverseContentBlock.html
-    # Note: text is handled as a primitive (string)
-}
-_BEDROCK_CONTENT_BLOCK_TYPES: set[_ContentBlockType] = set(_BEDROCK_CONTENT_BLOCK_FIELDS.keys())
-
-# Nested schemas for deep filtering of Bedrock content blocks
-_BEDROCK_CONTENT_BLOCK_SCHEMAS: dict[_ContentBlockType, dict[str, Any]] = {
-    "image": {
-        "format": True,
-        "source": {"bytes": True, "s3Location": {"bucket": True, "key": True, "region": True, "version": True}},
-    },
-    "toolResult": {"content": True, "toolUseId": True, "status": True},
-    "toolUse": {"input": True, "name": True, "toolUseId": True},
-    "document": {
-        "name": True,
-        "source": {"bytes": True, "s3Location": {"bucket": True, "key": True, "region": True, "version": True}},
-        "format": True,
-        "citations": True,
-        "context": True,
-    },
-    "video": {
-        "format": True,
-        "source": {"bytes": True, "s3Location": {"bucket": True, "key": True, "region": True, "version": True}},
-    },
-    "reasoningContent": {"reasoningText": {"text": True, "signature": True}, "redactedContent": True},
-    "citationsContent": {"citations": True, "content": True},
-    "cachePoint": {"type": True},
-    "guardContent": {
-        "image": {
-            "format": True,
-            "source": {"bytes": True, "s3Location": {"bucket": True, "key": True, "region": True, "version": True}},
-        },
-        "text": {"qualifiers": True, "text": True},
-    },
-}
-
 T = TypeVar("T", bound=BaseModel)
 
 DEFAULT_READ_TIMEOUT = 120
@@ -258,17 +181,6 @@ def get_config(self) -> BedrockConfig:
         """
         return self.config
 
-    def _should_include_tool_result_status(self) -> bool:
-        """Determine whether to include tool result status based on current config."""
-        include_status = self.config.get("include_tool_result_status", "auto")
-
-        if include_status is True:
-            return True
-        elif include_status is False:
-            return False
-        else:  # "auto"
-            return any(model in self.config["model_id"] for model in _MODELS_INCLUDE_STATUS)
-
     def format_request(
         self,
         messages: Messages,
@@ -352,7 +264,7 @@ def format_request(
             ),
         }
 
-    def _format_bedrock_messages(self, messages: Messages) -> Messages:
+    def _format_bedrock_messages(self, messages: Messages) -> list[dict[str, Any]]:
         """Format messages for Bedrock API compatibility.
 
         This function ensures messages conform to Bedrock's expected format by:
@@ -373,13 +285,13 @@ def _format_bedrock_messages(self, messages: Messages) -> Messages:
             content blocks to remove any additional fields before sending to Bedrock.
             https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlock.html
         """
-        cleaned_messages = []
+        cleaned_messages: list[dict[str, Any]] = []
 
         filtered_unknown_members = False
         dropped_deepseek_reasoning_content = False
 
         for message in messages:
-            cleaned_content: list[ContentBlock] = []
+            cleaned_content: list[dict[str, Any]] = []
 
             for content_block in message["content"]:
                 # Filter out SDK_UNKNOWN_MEMBER content blocks
@@ -393,28 +305,13 @@ def _format_bedrock_messages(self, messages: Messages) -> Messages:
                     dropped_deepseek_reasoning_content = True
                     continue
 
-                # Clean content blocks that need field filtering for Bedrock API compatibility
-                filterable_block_types = set(content_block.keys()) & _BEDROCK_CONTENT_BLOCK_TYPES
-
-                if filterable_block_types:
-                    # Should only be one block type per content block since it is a discriminated union
-                    block_type = cast(_ContentBlockType, next(iter(filterable_block_types)))
-                    block_data = content_block[block_type]
-                    schema = _BEDROCK_CONTENT_BLOCK_SCHEMAS[block_type].copy()
-
-                    if block_type == "toolResult" and not self._should_include_tool_result_status():
-                        schema.pop("status", None)
-
-                    cleaned_data = _deep_filter(block_data, schema)
-                    cleaned_content.append(cast(ContentBlock, {block_type: cleaned_data}))
-                else:
-                    # Keep other content blocks as-is
-                    cleaned_content.append(content_block)
+                # Format content blocks for Bedrock API compatibility
+                formatted_content = self._format_request_message_content(content_block)
+                cleaned_content.append(formatted_content)
 
             # Create new message with cleaned content (skip if empty)
             if cleaned_content:
-                cleaned_message: Message = Message(content=cleaned_content, role=message["role"])
-                cleaned_messages.append(cleaned_message)
+                cleaned_messages.append({"content": cleaned_content, "role": message["role"]})
 
         if filtered_unknown_members:
             logger.warning(
@@ -427,6 +324,184 @@ def _format_bedrock_messages(self, messages: Messages) -> Messages:
 
         return cleaned_messages
 
+    def _should_include_tool_result_status(self) -> bool:
+        """Determine whether to include tool result status based on current config."""
+        include_status = self.config.get("include_tool_result_status", "auto")
+
+        if include_status is True:
+            return True
+        elif include_status is False:
+            return False
+        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]:
+        """Format a Bedrock content block.
+
+        Bedrock strictly validates content blocks and throws exceptions for unknown fields.
+        This function extracts only the fields that Bedrock supports for each content type.
+
+        Args:
+            content: Content block to format.
+
+        Returns:
+            Bedrock formatted content block.
+
+        Raises:
+            TypeError: If the content block type is not supported by Bedrock.
+        """
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CachePointBlock.html
+        if "cachePoint" in content:
+            return {"cachePoint": {"type": content["cachePoint"]["type"]}}
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_DocumentBlock.html
+        if "document" in content:
+            document = content["document"]
+            result: dict[str, Any] = {}
+
+            # Handle required fields (all optional due to total=False)
+            if "name" in document:
+                result["name"] = document["name"]
+            if "format" in document:
+                result["format"] = document["format"]
+
+            # Handle source
+            if "source" in document:
+                result["source"] = {"bytes": document["source"]["bytes"]}
+
+            # Handle optional fields
+            if "citations" in document and document["citations"] is not None:
+                result["citations"] = {"enabled": document["citations"]["enabled"]}
+            if "context" in document:
+                result["context"] = document["context"]
+
+            return {"document": result}
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_GuardrailConverseContentBlock.html
+        if "guardContent" in content:
+            guard = content["guardContent"]
+            guard_text = guard["text"]
+            result = {"text": {"text": guard_text["text"], "qualifiers": guard_text["qualifiers"]}}
+            return {"guardContent": result}
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageBlock.html
+        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}
+            return {"image": result}
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ReasoningContentBlock.html
+        if "reasoningContent" in content:
+            reasoning = content["reasoningContent"]
+            result = {}
+
+            if "reasoningText" in reasoning:
+                reasoning_text = reasoning["reasoningText"]
+                result["reasoningText"] = {}
+                if "text" in reasoning_text:
+                    result["reasoningText"]["text"] = reasoning_text["text"]
+                # Only include signature if truthy (avoid empty strings)
+                if reasoning_text.get("signature"):
+                    result["reasoningText"]["signature"] = reasoning_text["signature"]
+
+            if "redactedContent" in reasoning:
+                result["redactedContent"] = reasoning["redactedContent"]
+
+            return {"reasoningContent": result}
+
+        # Pass through text and other simple content types
+        if "text" in content:
+            return {"text": content["text"]}
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolResultBlock.html
+        if "toolResult" in content:
+            tool_result = content["toolResult"]
+            formatted_content: list[dict[str, Any]] = []
+            for tool_result_content in tool_result["content"]:
+                if "json" in tool_result_content:
+                    # 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))
+                    )
+
+            result = {
+                "content": formatted_content,
+                "toolUseId": tool_result["toolUseId"],
+            }
+            if "status" in tool_result and self._should_include_tool_result_status():
+                result["status"] = tool_result["status"]
+            return {"toolResult": result}
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolUseBlock.html
+        if "toolUse" in content:
+            tool_use = content["toolUse"]
+            return {
+                "toolUse": {
+                    "input": tool_use["input"],
+                    "name": tool_use["name"],
+                    "toolUseId": tool_use["toolUseId"],
+                }
+            }
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_VideoBlock.html
+        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}
+            return {"video": result}
+
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CitationsContentBlock.html
+        if "citationsContent" in content:
+            citations = content["citationsContent"]
+            result = {}
+
+            if "citations" in citations:
+                result["citations"] = []
+                for citation in citations["citations"]:
+                    filtered_citation: dict[str, Any] = {}
+                    if "location" in citation:
+                        location = citation["location"]
+                        filtered_location = {}
+                        # Filter location fields to only include Bedrock-supported ones
+                        if "documentIndex" in location:
+                            filtered_location["documentIndex"] = location["documentIndex"]
+                        if "start" in location:
+                            filtered_location["start"] = location["start"]
+                        if "end" in location:
+                            filtered_location["end"] = location["end"]
+                        filtered_citation["location"] = filtered_location
+                    if "sourceContent" in citation:
+                        filtered_source_content: list[dict[str, Any]] = []
+                        for source_content in citation["sourceContent"]:
+                            if "text" in source_content:
+                                filtered_source_content.append({"text": source_content["text"]})
+                        if filtered_source_content:
+                            filtered_citation["sourceContent"] = filtered_source_content
+                    if "title" in citation:
+                        filtered_citation["title"] = citation["title"]
+                    result["citations"].append(filtered_citation)
+
+            if "content" in citations:
+                filtered_content: list[dict[str, Any]] = []
+                for generated_content in citations["content"]:
+                    if "text" in generated_content:
+                        filtered_content.append({"text": generated_content["text"]})
+                if filtered_content:
+                    result["content"] = filtered_content
+
+            return {"citationsContent": result}
+
+        raise TypeError(f"content_type=<{next(iter(content))}> | unsupported type")
+
     def _has_blocked_guardrail(self, guardrail_data: dict[str, Any]) -> bool:
         """Check if guardrail data contains any blocked policies.
 
@@ -836,34 +911,3 @@ async def structured_output(
             raise ValueError("No valid tool use or tool use input was found in the Bedrock response.")
 
         yield {"output": output_model(**output_response)}
-
-
-def _deep_filter(data: Union[dict[str, Any], Any], schema: dict[str, Any]) -> dict[str, Any]:
-    """Fast recursive filtering using nested dict schemas.
-
-    Args:
-        data: Input data to filter (content block or nested dict)
-        schema: Schema defining allowed fields and nested structure
-
-    Returns:
-        Filtered dictionary containing only schema-defined fields
-    """
-    if not isinstance(data, dict):
-        return {}
-
-    result = {}
-    for key in data.keys() & schema.keys():
-        value = data[key]
-        schema_spec = schema[key]
-
-        if schema_spec is True:
-            result[key] = value
-        elif isinstance(schema_spec, dict) and isinstance(value, dict):
-            filtered = _deep_filter(value, schema_spec)
-            if filtered:
-                result[key] = filtered
-        elif isinstance(schema_spec, dict) and isinstance(value, list):
-            result[key] = [_deep_filter(item, schema_spec) for item in value if isinstance(item, dict)]
-        else:
-            result[key] = value
-    return result