agentrun-sdk 0.1.2__py3-none-any.whl

agentrun_sdk/models/bedrock.py

@@ -0,0 +1,649 @@
+"""AWS Bedrock model provider.
+
+- Docs: https://aws.amazon.com/bedrock/
+"""
+
+import asyncio
+import json
+import logging
+import os
+from typing import Any, AsyncGenerator, Callable, Iterable, Literal, Optional, Type, TypeVar, Union
+
+import boto3
+from botocore.config import Config as BotocoreConfig
+from botocore.exceptions import ClientError
+from pydantic import BaseModel
+from typing_extensions import TypedDict, Unpack, override
+
+from ..event_loop import streaming
+from ..tools import convert_pydantic_to_tool_spec
+from ..types.content import ContentBlock, Message, Messages
+from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
+from ..types.streaming import StreamEvent
+from ..types.tools import ToolResult, ToolSpec
+from .model import Model
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_BEDROCK_MODEL_ID = "us.anthropic.claude-sonnet-4-20250514-v1:0"
+DEFAULT_BEDROCK_REGION = "us-west-2"
+
+BEDROCK_CONTEXT_WINDOW_OVERFLOW_MESSAGES = [
+    "Input is too long for requested model",
+    "input length and `max_tokens` exceed context limit",
+    "too many total text bytes",
+]
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class BedrockModel(Model):
+    """AWS Bedrock model provider implementation.
+
+    The implementation handles Bedrock-specific features such as:
+
+    - Tool configuration for function calling
+    - Guardrails integration
+    - Caching points for system prompts and tools
+    - Streaming responses
+    - Context window overflow detection
+    """
+
+    class BedrockConfig(TypedDict, total=False):
+        """Configuration options for Bedrock models.
+
+        Attributes:
+            additional_args: Any additional arguments to include in the request
+            additional_request_fields: Additional fields to include in the Bedrock request
+            additional_response_field_paths: Additional response field paths to extract
+            cache_prompt: Cache point type for the system prompt
+            cache_tools: Cache point type for tools
+            guardrail_id: ID of the guardrail to apply
+            guardrail_trace: Guardrail trace mode. Defaults to enabled.
+            guardrail_version: Version of the guardrail to apply
+            guardrail_stream_processing_mode: The guardrail processing mode
+            guardrail_redact_input: Flag to redact input if a guardrail is triggered. Defaults to True.
+            guardrail_redact_input_message: If a Bedrock Input guardrail triggers, replace the input with this message.
+            guardrail_redact_output: Flag to redact output if guardrail is triggered. Defaults to False.
+            guardrail_redact_output_message: If a Bedrock Output guardrail triggers, replace output with this message.
+            max_tokens: Maximum number of tokens to generate in the response
+            model_id: The Bedrock model ID (e.g., "us.anthropic.claude-sonnet-4-20250514-v1:0")
+            stop_sequences: List of sequences that will stop generation when encountered
+            streaming: Flag to enable/disable streaming. Defaults to True.
+            temperature: Controls randomness in generation (higher = more random)
+            top_p: Controls diversity via nucleus sampling (alternative to temperature)
+        """
+
+        additional_args: Optional[dict[str, Any]]
+        additional_request_fields: Optional[dict[str, Any]]
+        additional_response_field_paths: Optional[list[str]]
+        cache_prompt: Optional[str]
+        cache_tools: Optional[str]
+        guardrail_id: Optional[str]
+        guardrail_trace: Optional[Literal["enabled", "disabled", "enabled_full"]]
+        guardrail_stream_processing_mode: Optional[Literal["sync", "async"]]
+        guardrail_version: Optional[str]
+        guardrail_redact_input: Optional[bool]
+        guardrail_redact_input_message: Optional[str]
+        guardrail_redact_output: Optional[bool]
+        guardrail_redact_output_message: Optional[str]
+        max_tokens: Optional[int]
+        model_id: str
+        stop_sequences: Optional[list[str]]
+        streaming: Optional[bool]
+        temperature: Optional[float]
+        top_p: Optional[float]
+
+    def __init__(
+        self,
+        *,
+        boto_session: Optional[boto3.Session] = None,
+        boto_client_config: Optional[BotocoreConfig] = None,
+        region_name: Optional[str] = None,
+        **model_config: Unpack[BedrockConfig],
+    ):
+        """Initialize provider instance.
+
+        Args:
+            boto_session: Boto Session to use when calling the Bedrock Model.
+            boto_client_config: Configuration to use when creating the Bedrock-Runtime Boto Client.
+            region_name: AWS region to use for the Bedrock service.
+                Defaults to the AWS_REGION environment variable if set, or "us-west-2" if not set.
+            **model_config: Configuration options for the Bedrock model.
+        """
+        if region_name and boto_session:
+            raise ValueError("Cannot specify both `region_name` and `boto_session`.")
+
+        self.config = BedrockModel.BedrockConfig(model_id=DEFAULT_BEDROCK_MODEL_ID)
+        self.update_config(**model_config)
+
+        logger.debug("config=<%s> | initializing", self.config)
+
+        session = boto_session or boto3.Session()
+
+        # Add strands-agents to the request user agent
+        if boto_client_config:
+            existing_user_agent = getattr(boto_client_config, "user_agent_extra", None)
+
+            # Append 'strands-agents' to existing user_agent_extra or set it if not present
+            if existing_user_agent:
+                new_user_agent = f"{existing_user_agent} strands-agents"
+            else:
+                new_user_agent = "strands-agents"
+
+            client_config = boto_client_config.merge(BotocoreConfig(user_agent_extra=new_user_agent))
+        else:
+            client_config = BotocoreConfig(user_agent_extra="strands-agents")
+
+        resolved_region = region_name or session.region_name or os.environ.get("AWS_REGION") or DEFAULT_BEDROCK_REGION
+
+        self.client = session.client(
+            service_name="bedrock-runtime",
+            config=client_config,
+            region_name=resolved_region,
+        )
+
+        logger.debug("region=<%s> | bedrock client created", self.client.meta.region_name)
+
+    @override
+    def update_config(self, **model_config: Unpack[BedrockConfig]) -> None:  # type: ignore
+        """Update the Bedrock Model configuration with the provided arguments.
+
+        Args:
+            **model_config: Configuration overrides.
+        """
+        self.config.update(model_config)
+
+    @override
+    def get_config(self) -> BedrockConfig:
+        """Get the current Bedrock Model configuration.
+
+        Returns:
+            The Bedrock model configuration.
+        """
+        return self.config
+
+    def format_request(
+        self,
+        messages: Messages,
+        tool_specs: Optional[list[ToolSpec]] = None,
+        system_prompt: Optional[str] = None,
+    ) -> dict[str, Any]:
+        """Format a Bedrock converse stream request.
+
+        Args:
+            messages: List of message objects to be processed by the model.
+            tool_specs: List of tool specifications to make available to the model.
+            system_prompt: System prompt to provide context to the model.
+
+        Returns:
+            A Bedrock converse stream request.
+        """
+        return {
+            "modelId": self.config["model_id"],
+            "messages": self._format_bedrock_messages(messages),
+            "system": [
+                *([{"text": system_prompt}] if system_prompt else []),
+                *([{"cachePoint": {"type": self.config["cache_prompt"]}}] if self.config.get("cache_prompt") else []),
+            ],
+            **(
+                {
+                    "toolConfig": {
+                        "tools": [
+                            *[{"toolSpec": tool_spec} for tool_spec in tool_specs],
+                            *(
+                                [{"cachePoint": {"type": self.config["cache_tools"]}}]
+                                if self.config.get("cache_tools")
+                                else []
+                            ),
+                        ],
+                        "toolChoice": {"auto": {}},
+                    }
+                }
+                if tool_specs
+                else {}
+            ),
+            **(
+                {"additionalModelRequestFields": self.config["additional_request_fields"]}
+                if self.config.get("additional_request_fields")
+                else {}
+            ),
+            **(
+                {"additionalModelResponseFieldPaths": self.config["additional_response_field_paths"]}
+                if self.config.get("additional_response_field_paths")
+                else {}
+            ),
+            **(
+                {
+                    "guardrailConfig": {
+                        "guardrailIdentifier": self.config["guardrail_id"],
+                        "guardrailVersion": self.config["guardrail_version"],
+                        "trace": self.config.get("guardrail_trace", "enabled"),
+                        **(
+                            {"streamProcessingMode": self.config.get("guardrail_stream_processing_mode")}
+                            if self.config.get("guardrail_stream_processing_mode")
+                            else {}
+                        ),
+                    }
+                }
+                if self.config.get("guardrail_id") and self.config.get("guardrail_version")
+                else {}
+            ),
+            "inferenceConfig": {
+                key: value
+                for key, value in [
+                    ("maxTokens", self.config.get("max_tokens")),
+                    ("temperature", self.config.get("temperature")),
+                    ("topP", self.config.get("top_p")),
+                    ("stopSequences", self.config.get("stop_sequences")),
+                ]
+                if value is not None
+            },
+            **(
+                self.config["additional_args"]
+                if "additional_args" in self.config and self.config["additional_args"] is not None
+                else {}
+            ),
+        }
+
+    def _format_bedrock_messages(self, messages: Messages) -> Messages:
+        """Format messages for Bedrock API compatibility.
+
+        This function ensures messages conform to Bedrock's expected format by:
+        - Cleaning tool result content blocks by removing additional fields that may be
+          useful for retaining information in hooks but would cause Bedrock validation
+          exceptions when presented with unexpected fields
+        - Ensuring all message content blocks are properly formatted for the Bedrock API
+
+        Args:
+            messages: List of messages to format
+
+        Returns:
+            Messages formatted for Bedrock API compatibility
+
+        Note:
+            Bedrock will throw validation exceptions when presented with additional
+            unexpected fields in tool result blocks.
+            https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolResultBlock.html
+        """
+        cleaned_messages = []
+
+        for message in messages:
+            cleaned_content: list[ContentBlock] = []
+
+            for content_block in message["content"]:
+                if "toolResult" in content_block:
+                    # Create a new content block with only the cleaned toolResult
+                    tool_result: ToolResult = content_block["toolResult"]
+
+                    # Keep only the required fields for Bedrock
+                    cleaned_tool_result = ToolResult(
+                        content=tool_result["content"], toolUseId=tool_result["toolUseId"], status=tool_result["status"]
+                    )
+
+                    cleaned_block: ContentBlock = {"toolResult": cleaned_tool_result}
+                    cleaned_content.append(cleaned_block)
+                else:
+                    # Keep other content blocks as-is
+                    cleaned_content.append(content_block)
+
+            # Create new message with cleaned content
+            cleaned_message: Message = Message(content=cleaned_content, role=message["role"])
+            cleaned_messages.append(cleaned_message)
+
+        return cleaned_messages
+
+    def _has_blocked_guardrail(self, guardrail_data: dict[str, Any]) -> bool:
+        """Check if guardrail data contains any blocked policies.
+
+        Args:
+            guardrail_data: Guardrail data from trace information.
+
+        Returns:
+            True if any blocked guardrail is detected, False otherwise.
+        """
+        input_assessment = guardrail_data.get("inputAssessment", {})
+        output_assessments = guardrail_data.get("outputAssessments", {})
+
+        # Check input assessments
+        if any(self._find_detected_and_blocked_policy(assessment) for assessment in input_assessment.values()):
+            return True
+
+        # Check output assessments
+        if any(self._find_detected_and_blocked_policy(assessment) for assessment in output_assessments.values()):
+            return True
+
+        return False
+
+    def _generate_redaction_events(self) -> list[StreamEvent]:
+        """Generate redaction events based on configuration.
+
+        Returns:
+            List of redaction events to yield.
+        """
+        events: list[StreamEvent] = []
+
+        if self.config.get("guardrail_redact_input", True):
+            logger.debug("Redacting user input due to guardrail.")
+            events.append(
+                {
+                    "redactContent": {
+                        "redactUserContentMessage": self.config.get(
+                            "guardrail_redact_input_message", "[User input redacted.]"
+                        )
+                    }
+                }
+            )
+
+        if self.config.get("guardrail_redact_output", False):
+            logger.debug("Redacting assistant output due to guardrail.")
+            events.append(
+                {
+                    "redactContent": {
+                        "redactAssistantContentMessage": self.config.get(
+                            "guardrail_redact_output_message", "[Assistant output redacted.]"
+                        )
+                    }
+                }
+            )
+
+        return events
+
+    @override
+    async def stream(
+        self,
+        messages: Messages,
+        tool_specs: Optional[list[ToolSpec]] = None,
+        system_prompt: Optional[str] = None,
+        **kwargs: Any,
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """Stream conversation with the Bedrock model.
+
+        This method calls either the Bedrock converse_stream API or the converse API
+        based on the streaming parameter in the configuration.
+
+        Args:
+            messages: List of message objects to be processed by the model.
+            tool_specs: List of tool specifications to make available to the model.
+            system_prompt: System prompt to provide context to the model.
+            **kwargs: Additional keyword arguments for future extensibility.
+
+        Yields:
+            Model events.
+
+        Raises:
+            ContextWindowOverflowException: If the input exceeds the model's context window.
+            ModelThrottledException: If the model service is throttling requests.
+        """
+
+        def callback(event: Optional[StreamEvent] = None) -> None:
+            loop.call_soon_threadsafe(queue.put_nowait, event)
+            if event is None:
+                return
+
+        loop = asyncio.get_event_loop()
+        queue: asyncio.Queue[Optional[StreamEvent]] = asyncio.Queue()
+
+        thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt)
+        task = asyncio.create_task(thread)
+
+        while True:
+            event = await queue.get()
+            if event is None:
+                break
+
+            yield event
+
+        await task
+
+    def _stream(
+        self,
+        callback: Callable[..., None],
+        messages: Messages,
+        tool_specs: Optional[list[ToolSpec]] = None,
+        system_prompt: Optional[str] = None,
+    ) -> None:
+        """Stream conversation with the Bedrock model.
+
+        This method operates in a separate thread to avoid blocking the async event loop with the call to
+        Bedrock's converse_stream.
+
+        Args:
+            callback: Function to send events to the main thread.
+            messages: List of message objects to be processed by the model.
+            tool_specs: List of tool specifications to make available to the model.
+            system_prompt: System prompt to provide context to the model.
+
+        Raises:
+            ContextWindowOverflowException: If the input exceeds the model's context window.
+            ModelThrottledException: If the model service is throttling requests.
+        """
+        logger.debug("formatting request")
+        request = self.format_request(messages, tool_specs, system_prompt)
+        logger.debug("request=<%s>", request)
+
+        logger.debug("invoking model")
+        streaming = self.config.get("streaming", True)
+
+        try:
+            logger.debug("got response from model")
+            if streaming:
+                response = self.client.converse_stream(**request)
+                for chunk in response["stream"]:
+                    if (
+                        "metadata" in chunk
+                        and "trace" in chunk["metadata"]
+                        and "guardrail" in chunk["metadata"]["trace"]
+                    ):
+                        guardrail_data = chunk["metadata"]["trace"]["guardrail"]
+                        if self._has_blocked_guardrail(guardrail_data):
+                            for event in self._generate_redaction_events():
+                                callback(event)
+
+                    callback(chunk)
+
+            else:
+                response = self.client.converse(**request)
+                for event in self._convert_non_streaming_to_streaming(response):
+                    callback(event)
+
+                if (
+                    "trace" in response
+                    and "guardrail" in response["trace"]
+                    and self._has_blocked_guardrail(response["trace"]["guardrail"])
+                ):
+                    for event in self._generate_redaction_events():
+                        callback(event)
+
+        except ClientError as e:
+            error_message = str(e)
+
+            if e.response["Error"]["Code"] == "ThrottlingException":
+                raise ModelThrottledException(error_message) from e
+
+            if any(overflow_message in error_message for overflow_message in BEDROCK_CONTEXT_WINDOW_OVERFLOW_MESSAGES):
+                logger.warning("bedrock threw context window overflow error")
+                raise ContextWindowOverflowException(e) from e
+
+            region = self.client.meta.region_name
+
+            # add_note added in Python 3.11
+            if hasattr(e, "add_note"):
+                # Aid in debugging by adding more information
+                e.add_note(f"└ Bedrock region: {region}")
+                e.add_note(f"└ Model id: {self.config.get('model_id')}")
+
+                if (
+                    e.response["Error"]["Code"] == "AccessDeniedException"
+                    and "You don't have access to the model" in error_message
+                ):
+                    e.add_note(
+                        "└ For more information see "
+                        "https://strandsagents.com/latest/user-guide/concepts/model-providers/amazon-bedrock/#model-access-issue"
+                    )
+
+                if (
+                    e.response["Error"]["Code"] == "ValidationException"
+                    and "with on-demand throughput isn’t supported" in error_message
+                ):
+                    e.add_note(
+                        "└ For more information see "
+                        "https://strandsagents.com/latest/user-guide/concepts/model-providers/amazon-bedrock/#on-demand-throughput-isnt-supported"
+                    )
+
+            raise e
+
+        finally:
+            callback()
+            logger.debug("finished streaming response from model")
+
+    def _convert_non_streaming_to_streaming(self, response: dict[str, Any]) -> Iterable[StreamEvent]:
+        """Convert a non-streaming response to the streaming format.
+
+        Args:
+            response: The non-streaming response from the Bedrock model.
+
+        Returns:
+            An iterable of response events in the streaming format.
+        """
+        # Yield messageStart event
+        yield {"messageStart": {"role": response["output"]["message"]["role"]}}
+
+        # Process content blocks
+        for content in response["output"]["message"]["content"]:
+            # Yield contentBlockStart event if needed
+            if "toolUse" in content:
+                yield {
+                    "contentBlockStart": {
+                        "start": {
+                            "toolUse": {
+                                "toolUseId": content["toolUse"]["toolUseId"],
+                                "name": content["toolUse"]["name"],
+                            }
+                        },
+                    }
+                }
+
+                # For tool use, we need to yield the input as a delta
+                input_value = json.dumps(content["toolUse"]["input"])
+
+                yield {"contentBlockDelta": {"delta": {"toolUse": {"input": input_value}}}}
+            elif "text" in content:
+                # Then yield the text as a delta
+                yield {
+                    "contentBlockDelta": {
+                        "delta": {"text": content["text"]},
+                    }
+                }
+            elif "reasoningContent" in content:
+                # Then yield the reasoning content as a delta
+                yield {
+                    "contentBlockDelta": {
+                        "delta": {"reasoningContent": {"text": content["reasoningContent"]["reasoningText"]["text"]}}
+                    }
+                }
+
+                if "signature" in content["reasoningContent"]["reasoningText"]:
+                    yield {
+                        "contentBlockDelta": {
+                            "delta": {
+                                "reasoningContent": {
+                                    "signature": content["reasoningContent"]["reasoningText"]["signature"]
+                                }
+                            }
+                        }
+                    }
+
+            # Yield contentBlockStop event
+            yield {"contentBlockStop": {}}
+
+        # Yield messageStop event
+        yield {
+            "messageStop": {
+                "stopReason": response["stopReason"],
+                "additionalModelResponseFields": response.get("additionalModelResponseFields"),
+            }
+        }
+
+        # Yield metadata event
+        if "usage" in response or "metrics" in response or "trace" in response:
+            metadata: StreamEvent = {"metadata": {}}
+            if "usage" in response:
+                metadata["metadata"]["usage"] = response["usage"]
+            if "metrics" in response:
+                metadata["metadata"]["metrics"] = response["metrics"]
+            if "trace" in response:
+                metadata["metadata"]["trace"] = response["trace"]
+            yield metadata
+
+    def _find_detected_and_blocked_policy(self, input: Any) -> bool:
+        """Recursively checks if the assessment contains a detected and blocked guardrail.
+
+        Args:
+            input: The assessment to check.
+
+        Returns:
+            True if the input contains a detected and blocked guardrail, False otherwise.
+
+        """
+        # Check if input is a dictionary
+        if isinstance(input, dict):
+            # Check if current dictionary has action: BLOCKED and detected: true
+            if input.get("action") == "BLOCKED" and input.get("detected") and isinstance(input.get("detected"), bool):
+                return True
+
+            # Recursively check all values in the dictionary
+            for value in input.values():
+                if isinstance(value, dict):
+                    return self._find_detected_and_blocked_policy(value)
+                # Handle case where value is a list of dictionaries
+                elif isinstance(value, list):
+                    for item in value:
+                        return self._find_detected_and_blocked_policy(item)
+        elif isinstance(input, list):
+            # Handle case where input is a list of dictionaries
+            for item in input:
+                return self._find_detected_and_blocked_policy(item)
+        # Otherwise return False
+        return False
+
+    @override
+    async def structured_output(
+        self, output_model: Type[T], prompt: Messages, system_prompt: Optional[str] = None, **kwargs: Any
+    ) -> AsyncGenerator[dict[str, Union[T, Any]], None]:
+        """Get structured output from the model.
+
+        Args:
+            output_model: The output model to use for the agent.
+            prompt: The prompt messages to use for the agent.
+            system_prompt: System prompt to provide context to the model.
+            **kwargs: Additional keyword arguments for future extensibility.
+
+        Yields:
+            Model events with the last being the structured output.
+        """
+        tool_spec = convert_pydantic_to_tool_spec(output_model)
+
+        response = self.stream(messages=prompt, tool_specs=[tool_spec], system_prompt=system_prompt, **kwargs)
+        async for event in streaming.process_stream(response):
+            yield event
+
+        stop_reason, messages, _, _ = event["stop"]
+
+        if stop_reason != "tool_use":
+            raise ValueError(f'Model returned stop_reason: {stop_reason} instead of "tool_use".')
+
+        content = messages["content"]
+        output_response: dict[str, Any] | None = None
+        for block in content:
+            # if the tool use name doesn't match the tool spec name, skip, and if the block is not a tool use, skip.
+            # if the tool use name never matches, raise an error.
+            if block.get("toolUse") and block["toolUse"]["name"] == tool_spec["name"]:
+                output_response = block["toolUse"]["input"]
+            else:
+                continue
+
+        if output_response is None:
+            raise ValueError("No valid tool use or tool use input was found in the Bedrock response.")
+
+        yield {"output": output_model(**output_response)}