from __future__ import annotations import functools import typing import warnings from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterable, Iterator, Mapping, Sequence from contextlib import asynccontextmanager, contextmanager from contextvars import ContextVar from dataclasses import dataclass, field, replace from datetime import datetime, timedelta from itertools import count from threading import Lock from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, cast, overload from urllib.parse import parse_qs, urlparse import anyio.to_thread from pydantic_core import to_json from typing_extensions import ParamSpec, TypedDict, assert_never try: from botocore.client import BaseClient from botocore.exceptions import ( BotoCoreError, ClientError, ConnectionError as BotocoreConnectionError, HTTPClientError, ) from botocore.model import StructureShape except ImportError as _import_error: raise ImportError( 'Please install `boto3` to use the Bedrock model, ' 'you can use the `bedrock` optional group — `pip install "pydantic-ai-slim[bedrock]"`' ) from _import_error from pydantic_ai import ( AudioUrl, BinaryContent, CachePoint, CompactionPart, DocumentUrl, FilePart, FinishReason, ImageUrl, ModelMessage, ModelProfileSpec, ModelRequest, ModelResponse, ModelResponsePart, ModelResponseStreamEvent, NativeToolCallPart, NativeToolReturnPart, RetryPromptPart, SpeechPart, SystemPromptPart, TextContent, TextPart, ThinkingPart, ToolAvailabilityDeltaPart, ToolCallPart, ToolReturnPart, UploadedFile, UserPromptPart, VideoUrl, _utils, usage, ) from pydantic_ai._output import DEFAULT_OUTPUT_TOOL_NAME from pydantic_ai._run_context import RunContext from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError, UserError from pydantic_ai.messages import ( _tool_result_provenance_tags, # pyright: ignore[reportPrivateUsage] is_multi_modal_content, ) from pydantic_ai.models import ( Model, ModelRequestParameters, StreamedResponse, _suggest_known_model_id_from_provider_error, # pyright: ignore[reportPrivateUsage] _unconverted_speech_part_error, # pyright: ignore[reportPrivateUsage] _unsynthesized_tool_availability_delta_error, # pyright: ignore[reportPrivateUsage] check_allow_model_requests, download_item, ) from pydantic_ai.models._tool_choice import ResolvedToolChoice, resolve_tool_choice from pydantic_ai.native_tools import AbstractNativeTool, CodeExecutionTool from pydantic_ai.profiles import DEFAULT_THINKING_TAGS from pydantic_ai.profiles.anthropic import ( ANTHROPIC_SAMPLING_PARAMS, ANTHROPIC_THINKING_BUDGET_MAP, resolve_anthropic_effort, ) from pydantic_ai.profiles.openai import OPENAI_REASONING_EFFORT_MAP from pydantic_ai.providers import Provider, infer_provider from pydantic_ai.providers.bedrock import BedrockModelProfile, remove_bedrock_geo_prefix from pydantic_ai.settings import ModelSettings, ThinkingLevel, merge_model_settings from pydantic_ai.tools import ToolDefinition if TYPE_CHECKING: from botocore.eventstream import EventStream from mypy_boto3_bedrock_runtime import BedrockRuntimeClient from mypy_boto3_bedrock_runtime.literals import ( StopReasonType, ) from mypy_boto3_bedrock_runtime.type_defs import ( CachePointBlockTypeDef, ContentBlockOutputTypeDef, ContentBlockUnionTypeDef, ConverseRequestTypeDef, ConverseResponseTypeDef, ConverseStreamOutputTypeDef, ConverseStreamResponseTypeDef, ConverseTokensRequestTypeDef, CountTokensRequestTypeDef, DocumentSourceTypeDef, GuardrailConfigurationTypeDef, InferenceConfigurationTypeDef, JsonSchemaDefinitionTypeDef, MessageUnionTypeDef, OutputConfigTypeDef, PerformanceConfigurationTypeDef, PromptVariableValuesTypeDef, ReasoningContentBlockOutputTypeDef, S3LocationTypeDef, ServiceTierTypeDef, SystemContentBlockTypeDef, TokenUsageTypeDef, ToolChoiceTypeDef, ToolConfigurationTypeDef, ToolResultBlockOutputTypeDef, ToolResultContentBlockOutputTypeDef, ToolSpecificationTypeDef, ToolTypeDef, ToolUseBlockOutputTypeDef, ) @contextmanager def _map_api_errors(model_name: str, model_id_namespace: str = 'bedrock') -> Generator[None]: try: yield except ClientError as e: metadata = e.response.get('ResponseMetadata', {}) status_code = metadata.get('HTTPStatusCode') if isinstance(status_code, int): suggested_model_id = None error = e.response.get('Error') if _utils.is_str_dict(error) and error.get('Message') == 'The provided model identifier is invalid.': suggested_model_id = _suggest_known_model_id_from_provider_error(model_id_namespace, model_name) raise ModelHTTPError( status_code=status_code, model_name=model_name, body=e.response, headers=metadata.get('HTTPHeaders'), suggested_model_id=suggested_model_id, ) from e raise ModelAPIError(model_name=model_name, message=str(e)) from e except (HTTPClientError, BotocoreConnectionError) as e: # botocore raises transport failures (timeouts, connection errors) as `BotoCoreError`, not `ClientError`. raise ModelAPIError(model_name=model_name, message=str(e)) from e class _BotocoreRequestParams(TypedDict): headers: dict[str, str] @dataclass class _ExtraHeadersState: client_id: int headers: dict[str, str] claimed: bool = False _EXTRA_HEADERS_REGISTRATION_LOCK = Lock() _EXTRA_HEADERS_CONTEXT_KEY = 'pydantic_ai_extra_headers' _EXTRA_HEADERS_OPERATIONS = ('Converse', 'ConverseStream', 'CountTokens') _extra_headers_var: ContextVar[_ExtraHeadersState | None] = ContextVar('_extra_headers_var', default=None) _BedrockCallResult = TypeVar('_BedrockCallResult') def _claim_extra_headers(client_id: int, context: dict[str, Any], **_: Any) -> None: if (active := _extra_headers_var.get()) and active.client_id == client_id and not active.claimed: active.claimed = True context[_EXTRA_HEADERS_CONTEXT_KEY] = active.headers def _inject_extra_headers(params: _BotocoreRequestParams, context: dict[str, Any], **_: Any) -> None: extra_headers: dict[str, str] = context.pop(_EXTRA_HEADERS_CONTEXT_KEY, {}) headers = params['headers'] for key, value in extra_headers.items(): for existing_key in tuple(headers): if existing_key.lower() == key.lower(): del headers[existing_key] headers[key] = value def _register_extra_headers(client: BedrockRuntimeClient) -> None: """Register request-scoped header handlers once per client.""" # botocore's first registration mutates an unsynchronized handler trie and lookup cache; serialize it so a # concurrent pydantic-ai request can't emit against a half-updated cache. with _EXTRA_HEADERS_REGISTRATION_LOCK: for operation in _EXTRA_HEADERS_OPERATIONS: client.meta.events.register_first( f'provide-client-params.bedrock-runtime.{operation}', functools.partial(_claim_extra_headers, id(client)), unique_id=f'pydantic-ai-extra-headers-claim-{operation}', ) client.meta.events.register_first( f'before-call.bedrock-runtime.{operation}', _inject_extra_headers, unique_id=f'pydantic-ai-extra-headers-inject-{operation}', ) async def _call_bedrock( client: BedrockRuntimeClient, method: Callable[..., _BedrockCallResult], params: Mapping[str, Any], extra_headers: dict[str, str] | None, ) -> _BedrockCallResult: _register_extra_headers(client) headers = dict(extra_headers or {}) def call() -> _BedrockCallResult: context_token = _extra_headers_var.set(_ExtraHeadersState(id(client), headers)) try: return method(**params) finally: _extra_headers_var.reset(context_token) return await anyio.to_thread.run_sync(call) _SUPPORTED_IMAGE_FORMATS = ('jpeg', 'png', 'gif', 'webp') _SUPPORTED_VIDEO_FORMATS = ('mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp') _SUPPORTED_DOCUMENT_FORMATS = ('pdf', 'txt', 'csv', 'doc', 'docx', 'xls', 'xlsx', 'html', 'md') _ADAPTIVE_THINKING_SETTING = "bedrock_additional_model_requests_fields={'thinking': {'type': 'adaptive'}}" _BEDROCK_USAGE_FIELDS = frozenset( {'inputTokens', 'outputTokens', 'totalTokens', 'cacheReadInputTokens', 'cacheWriteInputTokens'} ) def _make_image_block(format: str, source: DocumentSourceTypeDef) -> ContentBlockUnionTypeDef: if format not in _SUPPORTED_IMAGE_FORMATS: raise UserError(f'Unsupported image format: {format}') return {'image': {'format': format, 'source': source}} def _make_video_block(format: str, source: DocumentSourceTypeDef) -> ContentBlockUnionTypeDef: if format not in _SUPPORTED_VIDEO_FORMATS: raise UserError(f'Unsupported video format: {format}') return {'video': {'format': format, 'source': source}} def _make_document_block(name: str, format: str, source: DocumentSourceTypeDef) -> ContentBlockUnionTypeDef: if format not in _SUPPORTED_DOCUMENT_FORMATS: raise UserError(f'Unsupported document format: {format}') return {'document': {'name': name, 'format': format, 'source': source}} # Content-block kinds that may appear in a user message alongside a `toolResult` block. Used as the # permissive default for `bedrock_tool_result_colocatable_content` (no model restriction). _ALL_TOOL_RESULT_COLOCATABLE_CONTENT: frozenset[Literal['text', 'image', 'document', 'video']] = frozenset( {'text', 'image', 'document', 'video'} ) LatestBedrockModelNames = Literal[ 'amazon.titan-tg1-large', 'amazon.titan-text-lite-v1', 'amazon.titan-text-express-v1', 'us.amazon.nova-2-lite-v1:0', 'us.amazon.nova-pro-v1:0', 'us.amazon.nova-lite-v1:0', 'us.amazon.nova-micro-v1:0', 'anthropic.claude-3-5-sonnet-20241022-v2:0', 'us.anthropic.claude-3-5-sonnet-20241022-v2:0', 'anthropic.claude-3-5-haiku-20241022-v1:0', 'us.anthropic.claude-3-5-haiku-20241022-v1:0', 'anthropic.claude-instant-v1', 'anthropic.claude-v2:1', 'anthropic.claude-v2', 'anthropic.claude-3-sonnet-20240229-v1:0', 'us.anthropic.claude-3-sonnet-20240229-v1:0', 'anthropic.claude-3-haiku-20240307-v1:0', 'us.anthropic.claude-3-haiku-20240307-v1:0', 'anthropic.claude-3-opus-20240229-v1:0', 'us.anthropic.claude-3-opus-20240229-v1:0', 'anthropic.claude-3-5-sonnet-20240620-v1:0', 'us.anthropic.claude-3-5-sonnet-20240620-v1:0', 'anthropic.claude-3-7-sonnet-20250219-v1:0', 'us.anthropic.claude-3-7-sonnet-20250219-v1:0', 'anthropic.claude-opus-4-20250514-v1:0', 'us.anthropic.claude-opus-4-20250514-v1:0', 'global.anthropic.claude-opus-4-5-20251101-v1:0', 'anthropic.claude-sonnet-4-20250514-v1:0', 'us.anthropic.claude-sonnet-4-20250514-v1:0', 'eu.anthropic.claude-sonnet-4-20250514-v1:0', 'anthropic.claude-sonnet-4-5-20250929-v1:0', 'us.anthropic.claude-sonnet-4-5-20250929-v1:0', 'eu.anthropic.claude-sonnet-4-5-20250929-v1:0', 'anthropic.claude-sonnet-4-6', 'us.anthropic.claude-sonnet-4-6', 'eu.anthropic.claude-sonnet-4-6', 'anthropic.claude-haiku-4-5-20251001-v1:0', 'us.anthropic.claude-haiku-4-5-20251001-v1:0', 'eu.anthropic.claude-haiku-4-5-20251001-v1:0', 'cohere.command-text-v14', 'cohere.command-r-v1:0', 'cohere.command-r-plus-v1:0', 'cohere.command-light-text-v14', 'meta.llama3-8b-instruct-v1:0', 'meta.llama3-70b-instruct-v1:0', 'meta.llama3-1-8b-instruct-v1:0', 'us.meta.llama3-1-8b-instruct-v1:0', 'meta.llama3-1-70b-instruct-v1:0', 'us.meta.llama3-1-70b-instruct-v1:0', 'meta.llama3-1-405b-instruct-v1:0', 'us.meta.llama3-2-11b-instruct-v1:0', 'us.meta.llama3-2-90b-instruct-v1:0', 'us.meta.llama3-2-1b-instruct-v1:0', 'us.meta.llama3-2-3b-instruct-v1:0', 'us.meta.llama3-3-70b-instruct-v1:0', 'mistral.mistral-7b-instruct-v0:2', 'mistral.mixtral-8x7b-instruct-v0:1', 'mistral.mistral-large-2402-v1:0', 'mistral.mistral-large-2407-v1:0', # Anthropic (models that require a cross-region inference profile) 'us.anthropic.claude-opus-4-1-20250805-v1:0', 'us.anthropic.claude-opus-4-5-20251101-v1:0', 'us.anthropic.claude-opus-4-6-v1', 'global.anthropic.claude-opus-4-6-v1', 'us.anthropic.claude-opus-4-7', 'global.anthropic.claude-opus-4-7', 'us.anthropic.claude-opus-4-8', 'global.anthropic.claude-opus-4-8', 'us.anthropic.claude-opus-5', 'global.anthropic.claude-opus-5', 'us.anthropic.claude-sonnet-5', 'global.anthropic.claude-sonnet-5', 'us.anthropic.claude-fable-5', 'us.anthropic.claude-fable-5-1', 'global.anthropic.claude-fable-5', 'global.anthropic.claude-fable-5-1', # Amazon Nova 'us.amazon.nova-premier-v1:0', 'global.amazon.nova-2-lite-v1:0', # OpenAI GPT-5.6 (models that require a cross-region inference profile) 'us.openai.gpt-5.6-sol', 'global.openai.gpt-5.6-sol', 'us.openai.gpt-5.6-luna', 'in.openai.gpt-5.6-luna', 'global.openai.gpt-5.6-luna', 'us.openai.gpt-5.6-terra', 'in.openai.gpt-5.6-terra', 'global.openai.gpt-5.6-terra', # Meta Llama 4 'us.meta.llama4-maverick-17b-instruct-v1:0', 'us.meta.llama4-scout-17b-instruct-v1:0', # Mistral 'mistral.mistral-small-2402-v1:0', 'mistral.mistral-large-3-675b-instruct', 'mistral.ministral-3-3b-instruct', 'mistral.ministral-3-8b-instruct', 'mistral.ministral-3-14b-instruct', 'mistral.magistral-small-2509', 'mistral.devstral-2-123b', 'mistral.pixtral-large-2502-v1:0', 'us.mistral.pixtral-large-2502-v1:0', # DeepSeek 'deepseek.r1-v1:0', 'deepseek.v3.2', # Qwen 'qwen.qwen3-32b-v1:0', 'qwen.qwen3-coder-30b-a3b-v1:0', 'qwen.qwen3-coder-next', 'qwen.qwen3-next-80b-a3b', 'qwen.qwen3-vl-235b-a22b', # Google Gemma 'google.gemma-3-4b-it', 'google.gemma-3-12b-it', 'google.gemma-3-27b-it', # MiniMax 'minimax.minimax-m2', 'minimax.minimax-m2.1', 'minimax.minimax-m2.5', # NVIDIA Nemotron 'nvidia.nemotron-nano-9b-v2', 'nvidia.nemotron-nano-12b-v2', 'nvidia.nemotron-nano-3-30b', 'nvidia.nemotron-super-3-120b', # Writer Palmyra (require a cross-region inference profile) 'us.writer.palmyra-x4-v1:0', 'us.writer.palmyra-x5-v1:0', # Z.AI GLM 'zai.glm-4.7', 'zai.glm-4.7-flash', 'zai.glm-5', # Moonshot AI Kimi 'moonshot.kimi-k2-thinking', 'moonshotai.kimi-k2.5', ] """Latest Bedrock models.""" BedrockModelName = str | LatestBedrockModelNames """Possible Bedrock model names. Since Bedrock supports a variety of date-stamped models, we explicitly list the latest models but allow any name in the type hints. See [the Bedrock docs](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) for a full list. """ P = ParamSpec('P') T = typing.TypeVar('T') _FINISH_REASON_MAP: dict[StopReasonType, FinishReason] = { 'content_filtered': 'content_filter', 'end_turn': 'stop', 'guardrail_intervened': 'content_filter', 'max_tokens': 'length', 'model_context_window_exceeded': 'length', 'stop_sequence': 'stop', 'malformed_model_output': 'error', 'malformed_tool_use': 'error', 'tool_use': 'tool_call', } def _parse_s3_source(url: str) -> DocumentSourceTypeDef: """Parse an S3 URL into a Bedrock DocumentSourceTypeDef.""" parsed = urlparse(url) s3_location: S3LocationTypeDef = {'uri': f'{parsed.scheme}://{parsed.netloc}{parsed.path}'} if bucket_owner := parse_qs(parsed.query).get('bucketOwner', [None])[0]: s3_location['bucketOwner'] = bucket_owner return {'s3Location': s3_location} def _insert_cache_point_before_trailing_documents( content: list[Any], cache_point: ContentBlockUnionTypeDef, *, raise_if_cannot_insert: bool = False, replace_existing: bool = False, ) -> bool: """Insert a cache point before trailing document/video content. AWS rejects cache points that directly follow documents and videos (but not images). This function finds the start of the trailing contiguous group of documents/videos and inserts a cache point before it. Args: content: The content list to modify in place. cache_point: The cache point block to insert. raise_if_cannot_insert: If True, raises UserError when cache point cannot be inserted (e.g., when the message contains only documents/videos). If False, silently skips. replace_existing: If True, an existing cache point at the insertion position is replaced (the latest marker wins). If False, it is kept and the new one skipped. Returns: True if a cache point was inserted, False otherwise. Raises: UserError: If raise_if_cannot_insert is True and the cache point cannot be placed. """ multimodal_keys = ['document', 'video'] # Find where the trailing contiguous group of documents/videos starts trailing_start: int | None = None for i in range(len(content) - 1, -1, -1): if any(key in content[i] for key in multimodal_keys): trailing_start = i else: break if trailing_start is not None and trailing_start > 0: prev_block = content[trailing_start - 1] if isinstance(prev_block, dict) and 'cachePoint' in prev_block: if replace_existing: content[trailing_start - 1] = cache_point return True return False content.insert(trailing_start, cache_point) return True elif trailing_start is None: # No trailing document/video content, append cache point at the end content.append(cache_point) return True else: # trailing_start == 0, can't insert at start if raise_if_cannot_insert: raise UserError( 'CachePoint cannot be placed when the user message contains only a document or video, ' 'due to Bedrock API restrictions. ' 'Add text content before or after your document or video to enable caching.' ) return False class BedrockModelSettings(ModelSettings, total=False): """Settings for Bedrock models. See [the Bedrock Converse API docs](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_RequestSyntax) for a full list. See [the boto3 implementation](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html) of the Bedrock Converse API. `extra_headers` are injected before the request is signed, so under SigV4 authentication they are covered by the signature (except the few headers botocore never signs, e.g. `X-Amzn-Trace-Id`). Headers the AWS SDK computes itself (e.g. `Authorization`, `User-Agent`, `X-Amz-Date`) are overwritten by botocore afterwards. """ # ALL FIELDS MUST BE `bedrock_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS. bedrock_guardrail_config: GuardrailConfigurationTypeDef """Content moderation and safety settings for Bedrock API requests. See more about it on . """ bedrock_performance_configuration: PerformanceConfigurationTypeDef """Performance optimization settings for model inference. See more about it on . """ bedrock_request_metadata: dict[str, str] """Additional metadata to attach to Bedrock API requests. See more about it on . """ bedrock_additional_model_response_fields_paths: list[str] """JSON paths to extract additional fields from model responses. See more about it on . """ bedrock_prompt_variables: Mapping[str, PromptVariableValuesTypeDef] """Variables for substitution into prompt templates. See more about it on . """ bedrock_additional_model_requests_fields: Mapping[str, Any] """Additional model-specific parameters to include in requests. See more about it on . """ bedrock_cache_tool_definitions: bool | Literal['5m', '1h'] """Whether to add a cache point after the last tool definition. When enabled, the last tool in the `tools` array will include a `cachePoint`, allowing Bedrock to cache tool definitions and reduce costs for compatible models. Set to `True` or `'5m'` for a 5-minute TTL (the default), or `'1h'` for a 1-hour TTL. See https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html for more information. """ bedrock_cache_instructions: bool | Literal['5m', '1h'] """Whether to add a cache point after the system prompt blocks. When enabled, an extra `cachePoint` is appended to the system prompt so Bedrock can cache system instructions. Set to `True` or `'5m'` for a 5-minute TTL (the default), or `'1h'` for a 1-hour TTL. See https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html for more information. """ bedrock_cache_messages: bool | Literal['5m', '1h'] """Convenience setting to enable caching for the last user message. When enabled, this automatically adds a cache point to the last content block in the final user message, which is useful for caching conversation history or context in multi-turn conversations. Set to `True` or `'5m'` for a 5-minute TTL (the default), or `'1h'` for a 1-hour TTL. Note: Uses 1 of Bedrock's 4 available cache points per request. Any additional CachePoint markers in messages will be automatically limited to respect the 4-cache-point maximum. See https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html for more information. """ bedrock_service_tier: ServiceTierTypeDef """Setting for optimizing performance and cost. Accepts `{'type': 'default' | 'flex' | 'priority' | 'reserved'}`. Takes precedence over the top-level [`service_tier`][pydantic_ai.settings.ModelSettings.service_tier], and is the only way to request `'reserved'` (which requires a pre-purchased capacity reservation). See more about it on . """ bedrock_inference_profile: str """An [inference profile](https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles.html) ARN to use as the `modelId` in API requests. When set, this value is used as the `modelId` in `converse` and `converse_stream` API calls instead of the base `model_name`. This allows you to pass the base model name (e.g. `'anthropic.claude-sonnet-4-5-20250929-v1:0'`) as `model_name` for detecting model capabilities and token counting, while routing requests through an inference profile for cost tracking or cross-region inference. """ @dataclass(init=False) class BedrockConverseModel(Model[BaseClient]): """A model that uses the Bedrock Converse API.""" _model_name: BedrockModelName = field(repr=False) _provider: Provider[BaseClient] = field(repr=False) _client: BaseClient | None = field(default=None, repr=False) def __init__( self, model_name: BedrockModelName, *, provider: Literal['bedrock', 'gateway'] | Provider[BaseClient] = 'bedrock', profile: ModelProfileSpec | None = None, settings: ModelSettings | None = None, ): """Initialize a Bedrock model. Args: model_name: The name of the model to use. model_name: The name of the Bedrock model to use. List of model names available [here](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html). provider: The provider to use for authentication and API access. Can be either the string 'bedrock' or an instance of `Provider[BaseClient]`. If not provided, a new provider will be created using the other parameters. profile: The model profile to use. Defaults to a profile picked by the provider based on the model name. settings: Model-specific settings that will be used as defaults for this model. """ self._model_name = model_name self._client = None if isinstance(provider, str): provider = infer_provider('gateway/bedrock' if provider == 'gateway' else provider) self._provider = provider super().__init__(settings=settings, profile=profile) if self.profile.get('bedrock_supported_on_converse', True) is False: raise UserError( f'Model {model_name!r} is not served by the Bedrock Converse API. Use `BedrockMantleProvider` ' "(the `bedrock-mantle:` prefix) to access it through Bedrock Mantle's OpenAI-compatible API." ) @property def client(self) -> BedrockRuntimeClient: """The boto3 client used to make requests to the Bedrock Converse API. Defaults to the client from the [`Provider`][pydantic_ai.providers.Provider]. It can be reassigned, e.g. to rotate short-lived credentials in a long-running service, but prefer assigning to [`BedrockProvider.client`][pydantic_ai.providers.bedrock.BedrockProvider.client] so all models sharing the provider pick up the new client. Once you've assigned a client here, you're responsible for keeping it valid; the provider's client is no longer consulted. """ return cast('BedrockRuntimeClient', self._client or self._provider.client) @client.setter def client(self, client: BedrockRuntimeClient) -> None: # Kept for backward compatibility (this used to be a plain attribute); `BedrockProvider.client` is the cleaner # place to swap the client, as it's shared by all models using the provider. self._client = client @property def base_url(self) -> str: return str(self.client.meta.endpoint_url) @property def model_name(self) -> str: """The model name.""" return self._model_name @property def system(self) -> str: """The model provider.""" return self._provider.name def resolve_prompt_cache_retention(self, model_settings: ModelSettings | None) -> timedelta | None: """Resolve the longest retention requested by supported Bedrock cache settings.""" settings = merge_model_settings(self.settings, model_settings) or {} return self._max_prompt_cache_retention( settings.get('bedrock_cache_instructions') if self.profile.get('bedrock_supports_prompt_caching', False) else None, settings.get('bedrock_cache_messages') if self.profile.get('bedrock_supports_prompt_caching', False) else None, settings.get('bedrock_cache_tool_definitions') if self.profile.get('bedrock_supports_tool_caching', False) else None, ) @classmethod def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]: """The set of builtin tool types this model can handle.""" return frozenset({CodeExecutionTool}) def prepare_request( self, model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters ) -> tuple[ModelSettings | None, ModelRequestParameters]: settings = merge_model_settings(self.settings, model_settings) profile = cast(BedrockModelProfile, self.profile) thinking_type = _effective_thinking_type(settings, model_request_parameters, profile) thinking_blocks_output_tools = _thinking_blocks_tool_forcing(thinking_type, profile) if model_request_parameters.output_tools and thinking_blocks_output_tools: supports_json_schema_output = self.profile.get('supports_json_schema_output', False) model_request_parameters = model_request_parameters.with_default_output_mode( 'native' if supports_json_schema_output else 'prompted' ) if ( model_request_parameters.output_mode == 'tool' and not model_request_parameters.allow_text_output ): # pragma: no branch # This would result in `toolChoice: any`, which isn't available here. suggested_output_type = 'NativeOutput' if supports_json_schema_output else 'PromptedOutput' remedy = f'Use `output_type={suggested_output_type}(...)` instead.' if thinking_type == 'adaptive': raise UserError( f'{self.model_name!r} does not support output tools when a thinking setting is ' f'configured, because it rejects the forced tool choice they require. {remedy}' ) if profile.get('bedrock_thinking_variant') != 'anthropic': raise UserError(f'Bedrock does not support thinking and output tools at the same time. {remedy}') if profile.get('bedrock_supports_adaptive_thinking', False) and _supports_tool_forcing(profile): remedy += f' Alternatively, `{_ADAPTIVE_THINKING_SETTING}` supports output tools.' raise UserError( f'Bedrock does not support extended thinking and output tools at the same time. {remedy}' ) # Resolve 'auto' to the profile default here (a no-op if already resolved above) so the # strict-forcing check below also applies when native mode is reached via the profile default # rather than an explicit `NativeOutput(...)`; `super().prepare_request()` would otherwise only # resolve it after `customize_request_parameters()` has already transformed the schema. model_request_parameters = model_request_parameters.with_default_output_mode( self.profile.get('default_structured_output_mode', 'tool') ) if ( self.profile.get('supports_json_schema_output', False) and model_request_parameters.output_mode == 'native' and model_request_parameters.output_object is not None ): # Bedrock's structured-output API requires `strict: true` on the output object — see # https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html # so we force it regardless of the caller's setting. Mirrors Anthropic's behavior. model_request_parameters = replace( model_request_parameters, output_object=replace(model_request_parameters.output_object, strict=True) ) # Pass unmerged model_settings; base class does its own merge prepared_settings, model_request_parameters = super().prepare_request(model_settings, model_request_parameters) if self.profile.get('anthropic_disallows_sampling_settings', False) and prepared_settings: filtered: ModelSettings = {**prepared_settings} self._drop_unsupported_sampling_settings(filtered) prepared_settings = filtered or None return prepared_settings, model_request_parameters def _drop_unsupported_sampling_settings(self, model_settings: ModelSettings) -> None: """Drop the sampling settings a flagged model rejects, warning like `AnthropicModel` does. `temperature` and `top_p` would reach `inferenceConfig` and unified `top_k` would reach `additionalModelRequestFields`; all three are rejected with a 400 by the models that set `anthropic_disallows_sampling_settings`. A user's own `bedrock_additional_model_requests_fields` is deliberately left alone: it is the raw escape hatch, so a value placed there is addressed to Bedrock directly. """ dropped = [setting for setting in ANTHROPIC_SAMPLING_PARAMS if setting in model_settings] for setting in dropped: model_settings.pop(setting, None) if dropped: warnings.warn( f'Sampling parameters {dropped} are not supported by {self.model_name!r}. These settings will be ignored.', UserWarning, stacklevel=2, ) @property def _botocore_supports_strict_tool_param(self) -> bool: """Whether the installed `botocore` knows the `strict` field on `toolSpec`. `botocore` validates request params against its own bundled service model, so a `botocore` older than the one that introduced strict tool calls rejects `strict` with a `ParamValidationError` regardless of what the Bedrock model itself supports. This notably happens on AWS Lambda, where the runtime's bundled `botocore` can shadow a newer one provided via a layer. """ tool_spec_shape = self.client.meta.service_model.shape_for('ToolSpecification') return isinstance(tool_spec_shape, StructureShape) and 'strict' in tool_spec_shape.members def _map_tool_definition(self, f: ToolDefinition) -> ToolTypeDef: tool_spec: ToolSpecificationTypeDef = {'name': f.name, 'inputSchema': {'json': f.parameters_json_schema}} if f.description: # pragma: no branch tool_spec['description'] = f.description if f.strict and self.profile.get('bedrock_supports_strict_tool_definition', False): if self._botocore_supports_strict_tool_param: tool_spec['strict'] = f.strict else: warnings.warn( 'The installed `botocore` is too old to send `strict` tool definitions to Bedrock, ' 'so the request is sent without `strict`. Upgrade `boto3`/`botocore` to enable strict ' "tool calls; on AWS Lambda, the runtime's bundled `botocore` may be shadowing a newer " 'one from your layer.', UserWarning, ) return {'toolSpec': tool_spec} @staticmethod def _native_output_format( model_request_parameters: ModelRequestParameters, ) -> OutputConfigTypeDef | None: """Build the `outputConfig` block for native structured output. See [Bedrock structured output](https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html). """ if model_request_parameters.output_mode != 'native' or model_request_parameters.output_object is None: return None output_object = model_request_parameters.output_object json_schema_config: JsonSchemaDefinitionTypeDef = { 'name': output_object.name or DEFAULT_OUTPUT_TOOL_NAME, 'schema': to_json(output_object.json_schema).decode(), } if output_object.description: json_schema_config['description'] = output_object.description return {'textFormat': {'type': 'json_schema', 'structure': {'jsonSchema': json_schema_config}}} async def request( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ModelResponse: check_allow_model_requests() model_settings, model_request_parameters = self.prepare_request( model_settings, model_request_parameters, ) settings = cast(BedrockModelSettings, model_settings or {}) response = await self._messages_create(messages, False, settings, model_request_parameters) model_response = await self._process_response(response) return model_response async def count_tokens( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> usage.RequestUsage: """Count the number of tokens, works with limited models. Check the actual supported models on """ check_allow_model_requests() model_settings, model_request_parameters = self.prepare_request(model_settings, model_request_parameters) settings = cast(BedrockModelSettings, model_settings or {}) system_prompt, bedrock_messages = await self._map_messages(messages, model_request_parameters, settings) converse: ConverseTokensRequestTypeDef = { 'messages': bedrock_messages, 'system': system_prompt, } # No native-tool strip is needed here (unlike Anthropic's count_tokens, which must drop server tools): # count-tokens-capable models (Claude) don't support native tools, and native-tool-capable models # (Nova-2) don't support count_tokens, so a `systemTool` can never reach this request. tool_config = self._map_tool_config(model_request_parameters, settings) if tool_config: converse['toolConfig'] = tool_config tools: list[ToolTypeDef] = list(tool_config['tools']) if tool_config else [] self._limit_cache_points(system_prompt, bedrock_messages, tools) if additional_model_requests_fields := self._build_additional_model_request_fields( settings, model_request_parameters ): converse['additionalModelRequestFields'] = additional_model_requests_fields params: CountTokensRequestTypeDef = { 'modelId': remove_bedrock_geo_prefix(self.model_name), 'input': {'converse': converse}, } # One client object for both registration and the call, in case the property is reassigned mid-request. client = self.client with _map_api_errors(self.model_name, self._provider.model_id_namespace): response = await _call_bedrock(client, client.count_tokens, params, settings.get('extra_headers')) return usage.RequestUsage(input_tokens=response['inputTokens']) @asynccontextmanager async def request_stream( self, messages: list[ModelMessage], model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, run_context: RunContext[Any] | None = None, ) -> AsyncGenerator[StreamedResponse]: check_allow_model_requests() model_settings, model_request_parameters = self.prepare_request( model_settings, model_request_parameters, ) settings = cast(BedrockModelSettings, model_settings or {}) response = await self._messages_create(messages, True, settings, model_request_parameters) yield BedrockStreamedResponse( model_request_parameters=model_request_parameters, _model_name=self.model_name, _model_profile=cast(BedrockModelProfile, self.profile), _event_stream=response['stream'], _provider_name=self._provider.name, _model_id_namespace=self._provider.model_id_namespace, _provider_url=self.base_url, _provider_response_id=response.get('ResponseMetadata', {}).get('RequestId', None), ) async def _process_response(self, response: ConverseResponseTypeDef) -> ModelResponse: items: list[ModelResponsePart] = [] if message := response['output'].get('message'): # pragma: no branch for item in message['content']: if reasoning_content := item.get('reasoningContent'): if redacted_content := reasoning_content.get('redactedContent'): items.append( ThinkingPart( id='redacted_content', content='', signature=redacted_content.decode('utf-8'), provider_name=self.system, ) ) elif reasoning_text := reasoning_content.get('reasoningText'): # pragma: no branch signature = reasoning_text.get('signature') items.append( ThinkingPart( content=reasoning_text['text'], signature=signature, provider_name=self.system if signature else None, ) ) if text := item.get('text'): items.append(TextPart(content=text)) elif tool_use := item.get('toolUse'): if tool_use.get('type') == 'server_tool_use': if tool_use['name'] == 'nova_code_interpreter': # pragma: no branch call_part = NativeToolCallPart( provider_name=self.system, tool_name=CodeExecutionTool.kind, args=tool_use['input'], tool_call_id=tool_use['toolUseId'], ) call_part.otel_metadata = {'code_arg_name': 'snippet', 'code_arg_language': 'python'} items.append(call_part) else: items.append( ToolCallPart( tool_name=tool_use['name'], args=tool_use['input'], tool_call_id=tool_use['toolUseId'], ), ) elif tool_result := item.get('toolResult'): if tool_result.get('type') == 'nova_code_interpreter_result': # pragma: no branch items.append( NativeToolReturnPart( provider_name=self.system, tool_name=CodeExecutionTool.kind, content=tool_result['content'][0].get('json') if tool_result['content'] else None, tool_call_id=tool_result.get('toolUseId'), provider_details={'status': tool_result['status']} if 'status' in tool_result else {}, ) ) u = _map_usage(response['usage'], self._provider.name, self.base_url, self.model_name) response_id = response.get('ResponseMetadata', {}).get('RequestId', None) raw_finish_reason = response['stopReason'] provider_details: dict[str, Any] = {'finish_reason': raw_finish_reason} if 'trace' in response: provider_details['trace'] = response['trace'] finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason) return ModelResponse( parts=items, usage=u, model_name=self.model_name, provider_response_id=response_id, provider_name=self._provider.name, provider_url=self.base_url, finish_reason=finish_reason, provider_details=provider_details, ) def _build_additional_model_request_fields( self, model_settings: BedrockModelSettings, model_request_parameters: ModelRequestParameters, ) -> dict[str, Any] | None: """Build `additionalModelRequestFields` from user-supplied fields plus unified `top_k` and `thinking`.""" existing = dict(model_settings.get('bedrock_additional_model_requests_fields') or {}) profile = self.profile # Bedrock's `inferenceConfig` has no `topK`, so unified `top_k` rides in the model-specific # `additionalModelRequestFields` (shape varies per family). A user-supplied key wins. if (top_k := model_settings.get('top_k')) is not None: if profile.get('bedrock_top_k_variant', None) == 'anthropic' and 'top_k' not in existing: existing['top_k'] = top_k elif profile.get('bedrock_top_k_variant', None) == 'nova': # Nova nests `topK` under `inferenceConfig`, so check that specific key (not the parent) # and merge into a fresh dict to preserve any other user-supplied `inferenceConfig` fields # without mutating the user's settings in place. A user-supplied `topK` wins. inference_config: Mapping[str, Any] = existing.get('inferenceConfig') or {} if isinstance(inference_config, dict) and 'topK' not in inference_config: existing['inferenceConfig'] = {**inference_config, 'topK': top_k} thinking = model_request_parameters.thinking if thinking is None: return existing or None variant = profile.get('bedrock_thinking_variant', None) if variant == 'anthropic' and 'thinking' not in existing: if profile.get('bedrock_supports_adaptive_thinking', False): if thinking is not False: existing['thinking'] = {'type': 'adaptive'} # Bedrock puts effort in output_config (a sibling of thinking), matching the direct Anthropic API shape. # The merged profile carries the Anthropic xhigh flag, so `xhigh` passes through on the same # models as the direct API. Bedrock rejects it on the other models, where it maps to `max`. if ( profile.get('bedrock_supports_effort', False) and isinstance(thinking, str) and 'output_config' not in existing ): effort = resolve_anthropic_effort( thinking, supports_xhigh=profile.get('anthropic_supports_xhigh_effort', False) ) existing['output_config'] = {'effort': effort} elif thinking is False: existing['thinking'] = {'type': 'disabled'} else: existing['thinking'] = {'type': 'enabled', 'budget_tokens': ANTHROPIC_THINKING_BUDGET_MAP[thinking]} elif variant == 'openai' and 'reasoning_effort' not in existing: if thinking is not False: # Bedrock doesn't accept reasoning_effort='none' existing['reasoning_effort'] = OPENAI_REASONING_EFFORT_MAP[thinking] elif variant == 'qwen' and 'reasoning_config' not in existing: if thinking is not False: # Qwen only supports low/high; map others to closest level_map: dict[ThinkingLevel, str] = { True: 'high', 'minimal': 'low', 'low': 'low', 'medium': 'high', 'high': 'high', 'xhigh': 'high', } existing['reasoning_config'] = level_map[thinking] return existing or None @overload async def _messages_create( self, messages: list[ModelMessage], stream: Literal[True], model_settings: BedrockModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ConverseStreamResponseTypeDef: pass @overload async def _messages_create( self, messages: list[ModelMessage], stream: Literal[False], model_settings: BedrockModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ConverseResponseTypeDef: pass async def _messages_create( self, messages: list[ModelMessage], stream: bool, model_settings: BedrockModelSettings | None, model_request_parameters: ModelRequestParameters, ) -> ConverseResponseTypeDef | ConverseStreamResponseTypeDef: settings = model_settings or BedrockModelSettings() system_prompt, bedrock_messages = await self._map_messages(messages, model_request_parameters, settings) inference_config = self._map_inference_config(settings) params: ConverseRequestTypeDef = { 'modelId': settings.get('bedrock_inference_profile') or self.model_name, 'messages': bedrock_messages, 'system': system_prompt, 'inferenceConfig': inference_config, } tool_config = self._map_tool_config(model_request_parameters, settings) if tool_config: params['toolConfig'] = tool_config tools: list[ToolTypeDef] = list(tool_config['tools']) if tool_config else [] self._limit_cache_points(system_prompt, bedrock_messages, tools) if output_config := self._native_output_format(model_request_parameters): params['outputConfig'] = output_config # Bedrock supports a set of specific extra parameters if model_settings: if guardrail_config := model_settings.get('bedrock_guardrail_config', None): params['guardrailConfig'] = guardrail_config if performance_configuration := model_settings.get('bedrock_performance_configuration', None): params['performanceConfig'] = performance_configuration if request_metadata := model_settings.get('bedrock_request_metadata', None): params['requestMetadata'] = request_metadata if additional_model_response_fields_paths := model_settings.get( 'bedrock_additional_model_response_fields_paths', None ): params['additionalModelResponseFieldPaths'] = additional_model_response_fields_paths if prompt_variables := model_settings.get('bedrock_prompt_variables', None): params['promptVariables'] = prompt_variables if service_tier := model_settings.get('bedrock_service_tier'): params['serviceTier'] = service_tier elif (unified_tier := model_settings.get('service_tier')) and unified_tier != 'auto': params['serviceTier'] = {'type': unified_tier} if additional_model_requests_fields := self._build_additional_model_request_fields( settings, model_request_parameters ): params['additionalModelRequestFields'] = additional_model_requests_fields # One client object for both registration and the call, in case the property is reassigned mid-request. client = self.client with _map_api_errors(self.model_name, self._provider.model_id_namespace): if stream: model_response = await _call_bedrock( client, client.converse_stream, params, settings.get('extra_headers') ) else: model_response = await _call_bedrock(client, client.converse, params, settings.get('extra_headers')) return model_response @staticmethod def _map_inference_config( model_settings: ModelSettings | None, ) -> InferenceConfigurationTypeDef: model_settings = model_settings or {} inference_config: InferenceConfigurationTypeDef = {} if max_tokens := model_settings.get('max_tokens'): inference_config['maxTokens'] = max_tokens if (temperature := model_settings.get('temperature')) is not None: inference_config['temperature'] = temperature if (top_p := model_settings.get('top_p')) is not None: inference_config['topP'] = top_p if stop_sequences := model_settings.get('stop_sequences'): inference_config['stopSequences'] = stop_sequences return inference_config def _map_tool_config( self, model_request_parameters: ModelRequestParameters, model_settings: BedrockModelSettings | None, ) -> ToolConfigurationTypeDef | None: resolved_tool_choice = resolve_tool_choice(model_settings, model_request_parameters) tool_defs = model_request_parameters.declared_tool_defs profile = cast(BedrockModelProfile, self.profile) supports = _support_tool_forcing( self.model_name, profile, model_settings, model_request_parameters, resolved_tool_choice ) tool_choice: ToolChoiceTypeDef if resolved_tool_choice == 'auto': tool_choice = {'auto': {}} elif resolved_tool_choice == 'required': tool_choice = {'any': {}} if supports else {'auto': {}} elif resolved_tool_choice == 'none': # Bedrock doesn't support a native 'none' mode, so we don't send tools # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html return None elif isinstance(resolved_tool_choice, tuple): tool_choice_mode, tool_names = resolved_tool_choice if tool_choice_mode == 'required' and len(tool_names) == 1: if supports: tool_choice = {'tool': {'name': next(iter(tool_names))}} else: # Breaks caching, but native `toolChoice.tool` is unavailable here (unsupported profile or thinking enabled) tool_defs = {k: v for k, v in tool_defs.items() if k in tool_names} tool_choice = {'auto': {}} else: # Breaks caching, but Bedrock's toolChoice only supports a single tool name tool_defs = {k: v for k, v in tool_defs.items() if k in tool_names} tool_choice = {'auto': {}} if tool_choice_mode == 'auto' or not supports else {'any': {}} else: assert_never(resolved_tool_choice) tools: list[ToolTypeDef] = [self._map_tool_definition(t) for t in tool_defs.values()] for tool in model_request_parameters.native_tools: if tool.kind == CodeExecutionTool.kind: tools.append({'systemTool': {'name': 'nova_code_interpreter'}}) else: raise NotImplementedError( f"Native tool '{tool.kind}' is not supported yet. If it should be, please file an issue." ) if not tools: return None if cache_tool_definitions := (model_settings or {}).get('bedrock_cache_tool_definitions'): if profile.get('bedrock_supports_tool_caching', False): tools.append(cast('ToolTypeDef', self._get_cache_point(cache_tool_definitions))) tool_config: ToolConfigurationTypeDef = {'tools': tools} if tool_choice and profile.get('bedrock_supports_tool_choice', False): tool_config['toolChoice'] = tool_choice return tool_config async def _map_messages( # noqa: C901 self, messages: Sequence[ModelMessage], model_request_parameters: ModelRequestParameters, model_settings: BedrockModelSettings | None, ) -> tuple[list[SystemContentBlockTypeDef], list[MessageUnionTypeDef]]: """Maps a `pydantic_ai.Message` to the Bedrock `MessageUnionTypeDef`. Groups consecutive ToolReturnPart objects into a single user message as required by Bedrock Claude/Nova models. """ settings = model_settings or BedrockModelSettings() profile = self.profile system_prompt: list[SystemContentBlockTypeDef] = [] bedrock_messages: list[MessageUnionTypeDef] = [] document_count: Iterator[int] = count(1) # Content-block kinds that may share a user turn with a `toolResult` block for this model. colocatable_content = profile.get( 'bedrock_tool_result_colocatable_content', _ALL_TOOL_RESULT_COLOCATABLE_CONTENT ) # Most families accept a `status` field on `toolResult` blocks; Writer Palmyra rejects it. supports_tool_result_status = profile.get('bedrock_supports_tool_result_status', True) # Media returned from a tool that can't live inside a `toolResult` block (see # `bedrock_supported_media_kinds_in_tool_returns`) is emitted as a sibling block. Models like # Mistral and Llama require every `toolResult` for a tool-use turn to sit together in the message # immediately following it, with nothing else sharing that turn, so such sibling media can't be # placed there. When the media kind can't co-locate with a `toolResult` for this model, we collect # it across the whole consecutive tool-return group and flush it as a separate user message after # the grouped tool results; the merge pass below then separates it with a synthetic assistant turn. # Media that this model does allow alongside a `toolResult` stays co-located in the same turn. deferred_media_content: list[ContentBlockUnionTypeDef] = [] def flush_deferred_media() -> None: if deferred_media_content: bedrock_messages.append({'role': 'user', 'content': [*deferred_media_content]}) deferred_media_content.clear() for message in messages: if isinstance(message, ModelRequest): for part in message.parts: if isinstance(part, SystemPromptPart): if part.content: # pragma: no branch system_prompt.append({'text': part.content}) elif isinstance(part, UserPromptPart): flush_deferred_media() bedrock_messages.extend( await self._map_user_prompt( part, document_count, supports_prompt_caching=profile.get('bedrock_supports_prompt_caching', False), prior_messages=bedrock_messages, ) ) elif isinstance(part, ToolReturnPart): assert part.tool_call_id is not None tool_result_content: list[Any] = [] colocated_media_content: list[ContentBlockUnionTypeDef] = [] content_mode: Literal['str', 'jsonable'] = ( 'str' if profile.get('bedrock_tool_result_format', 'text') == 'text' else 'jsonable' ) # Two mutually exclusive ways to render a failed return, picked here so the loop # below stays free of per-item failure guards: # - No native error status: fold the failure into one wrapped `{'error': ...}` text # block, then iterate only the files. Each file still gets its "See file X." # reference below so the model can cross-reference the media with the result. # - Otherwise (success, or failed with `status='error'` set below): send every # content item verbatim; the status field carries the failure signal unwrapped. items: Sequence[Any] if part.outcome == 'failed' and not supports_tool_result_status: tool_result_content.append({'text': part.model_response_str()}) items = part.files else: items = part.content_items(mode=content_mode, wrap_if_error=False) for item in items: if isinstance(item, UploadedFile): self._validate_uploaded_file_provider(item) if not item.file_id.startswith('s3://'): raise UserError( f'UploadedFile for Bedrock must use an S3 URL (s3://bucket/key), got: {item.file_id}' ) uf_source = _parse_s3_source(item.file_id) try: uf_format = item.format except ValueError as e: raise UserError( f'Unsupported media type for Bedrock UploadedFile: {item.media_type}' ) from e if item.media_type.startswith('image/'): tool_result_content.append(_make_image_block(uf_format, uf_source)) elif item.media_type.startswith('video/'): tool_result_content.append(_make_video_block(uf_format, uf_source)) elif item.media_type.startswith('audio/'): raise UserError('Audio files are not supported for Bedrock UploadedFile') else: tool_result_content.append( _make_document_block(f'Document {next(document_count)}', uf_format, uf_source) ) elif is_multi_modal_content(item): if isinstance(item, AudioUrl): raise NotImplementedError('AudioUrl is not supported in Bedrock tool returns') file_block = await self._map_file_to_content_block(item, document_count) # pyright: ignore[reportArgumentType] kind = next((k for k in ('image', 'document', 'video') if k in file_block), None) if kind in profile.get( 'bedrock_supported_media_kinds_in_tool_returns', frozenset({'image'}) ): tool_result_content.append(file_block) else: tool_result_content.append({'text': f'See file {item.identifier}.'}) # This media lands on a user turn, so it is framed with the call it # came from — see `_tool_result_provenance_tags`. open_tag, close_tag = _tool_result_provenance_tags( part.tool_name, part.tool_call_id, item.identifier ) framed_media: list[ContentBlockUnionTypeDef] = [ {'text': open_tag}, file_block, {'text': close_tag}, ] if kind in colocatable_content: # This model allows the media alongside the `toolResult`; keep it in the same turn. colocated_media_content.extend(framed_media) else: # The media can't share the `toolResult`'s turn; defer it to a later user turn. deferred_media_content.extend(framed_media) else: tool_result_content.append({'text': item} if isinstance(item, str) else {'json': item}) if not tool_result_content: tool_result_content.append( {'text': str(part.content)} if content_mode == 'str' else {'json': part.content} ) success_result: ToolResultBlockOutputTypeDef = { 'toolUseId': part.tool_call_id, 'content': tool_result_content, } if supports_tool_result_status: success_result['status'] = 'error' if part.outcome == 'failed' else 'success' bedrock_messages.append( { 'role': 'user', 'content': [{'toolResult': success_result}, *colocated_media_content], } ) elif isinstance(part, RetryPromptPart): if part.tool_name is None: flush_deferred_media() bedrock_messages.append({'role': 'user', 'content': [{'text': part.model_response()}]}) else: assert part.tool_call_id is not None error_result: ToolResultBlockOutputTypeDef = { 'toolUseId': part.tool_call_id, 'content': [{'text': part.model_response()}], } if supports_tool_result_status: error_result['status'] = 'error' bedrock_messages.append({'role': 'user', 'content': [{'toolResult': error_result}]}) elif isinstance(part, ToolAvailabilityDeltaPart): # pragma: no cover raise _unsynthesized_tool_availability_delta_error() elif isinstance(part, SpeechPart): # pragma: no cover # Unconverted realtime speech; `prepare_messages` turns these into `UserPromptPart`s in `Model.prepare_messages`. raise _unconverted_speech_part_error() else: assert_never(part) elif isinstance(message, ModelResponse): flush_deferred_media() content: list[ContentBlockOutputTypeDef] = [] for item in message.parts: if isinstance(item, TextPart): content.append({'text': item.content}) elif isinstance(item, ThinkingPart): if ( item.provider_name == self.system and item.signature and profile.get('bedrock_send_back_thinking_parts', False) ): reasoning_content: ReasoningContentBlockOutputTypeDef if item.id == 'redacted_content': reasoning_content = { 'redactedContent': item.signature.encode('utf-8'), } else: reasoning_content = { 'reasoningText': { 'text': item.content, 'signature': item.signature, } } content.append({'reasoningContent': reasoning_content}) else: start_tag, end_tag = self.profile.get('thinking_tags', DEFAULT_THINKING_TAGS) content.append({'text': '\n'.join([start_tag, item.content, end_tag])}) elif isinstance(item, NativeToolCallPart): if item.provider_name == self.system: if item.tool_name == CodeExecutionTool.kind: server_tool_use_block_param: ToolUseBlockOutputTypeDef = { 'toolUseId': _utils.guard_tool_call_id(t=item), 'name': 'nova_code_interpreter', 'input': item.args_as_dict(), 'type': 'server_tool_use', } content.append({'toolUse': server_tool_use_block_param}) elif isinstance(item, NativeToolReturnPart): if item.provider_name == self.system: if item.tool_name == CodeExecutionTool.kind: result_content: list[ToolResultContentBlockOutputTypeDef] = [ {'json': cast(dict[str, Any], item.content)} ] tool_result: ToolResultBlockOutputTypeDef = { 'toolUseId': _utils.guard_tool_call_id(t=item), 'content': result_content, 'type': 'nova_code_interpreter_result', } if item.provider_details and 'status' in item.provider_details: tool_result['status'] = item.provider_details['status'] content.append({'toolResult': tool_result}) elif isinstance(item, CompactionPart | FilePart): # Compaction and file parts are not sent back to models that don't support them. pass # pragma: no cover else: assert isinstance(item, ToolCallPart) content.append(self._map_tool_call(item)) if content: bedrock_messages.append({'role': 'assistant', 'content': content}) else: assert_never(message) # Flush any tool-return media that trails the conversation (the common case: history ends with # tool returns and no following assistant turn). flush_deferred_media() # Merge together sequential user messages. Some models reject a user message that co-locates a # `toolResult` block with other content: Anthropic rejects documents and video next to it, while # Llama and Mistral reject anything sharing the turn (the `toolResult` must be alone). When the # combined content isn't co-locatable (per `colocatable_content`), split the turns instead of # merging. See https://github.com/pydantic/pydantic-ai/issues/6081 and `bedrock_tool_result_colocatable_content`. # `cachePoint` blocks are cache markers, not content, so they never force a split. processed_messages: list[MessageUnionTypeDef] = [] last_message: dict[str, Any] | None = None for current_message in bedrock_messages: if ( last_message is not None and current_message['role'] == last_message['role'] and current_message['role'] == 'user' ): merged_content = [*last_message['content'], *current_message['content']] has_tool_result = any('toolResult' in block for block in merged_content) has_non_colocatable = any( 'toolResult' not in block and 'cachePoint' not in block and next(iter(block)) not in colocatable_content for block in merged_content ) if has_tool_result and has_non_colocatable: # The `toolResult` can't share this model's turn with the other content. Bedrock # re-merges consecutive same-role turns, so a bare split isn't enough; separate the # two user turns with a synthetic assistant turn. Several models reject whitespace-only # text, so use a period. processed_messages.append({'role': 'assistant', 'content': [{'text': '.'}]}) else: # Add the new user content onto the existing user message. last_content = list(last_message['content']) last_content.extend(current_message['content']) last_message['content'] = last_content continue # Add the entire message to the list of messages. processed_messages.append(current_message) last_message = cast(dict[str, Any], current_message) if instruction_parts := self._get_instruction_parts(messages, model_request_parameters): for part in instruction_parts: system_prompt.append({'text': part.content}) if ( system_prompt and (cache_instructions := settings.get('bedrock_cache_instructions')) and profile.get('bedrock_supports_prompt_caching', False) ): cache_point = cast('SystemContentBlockTypeDef', self._get_cache_point(cache_instructions)) if instruction_parts and any(p.dynamic for p in instruction_parts): # Insert cache point after the last static instruction (static parts are sorted first) num_pre_instruction_blocks = len(system_prompt) - len(instruction_parts) num_static = sum(1 for p in instruction_parts if not p.dynamic) cache_idx = num_pre_instruction_blocks + num_static if cache_idx > 0: system_prompt.insert(cache_idx, cache_point) else: # All static or no instruction_parts: cache point at end. system_prompt.append(cache_point) if processed_messages and (cache_messages := settings.get('bedrock_cache_messages')): if profile.get('bedrock_supports_prompt_caching', False): last_user_content = self._get_last_user_message_content(processed_messages) if last_user_content is not None: # Note: `_get_last_user_message_content` ensures content doesn't already end with a `cachePoint`. _insert_cache_point_before_trailing_documents( last_user_content, self._get_cache_point(cache_messages) ) # Bedrock's Converse API requires at least one message, so an empty conversation (only a # system prompt/instructions) always needs a synthetic user turn. Beyond that, most model # families also reject a conversation that starts with an assistant turn (e.g. a # `message_history` that begins with a `ModelResponse`) with "A conversation must start with # a user message...", so we synthesize a leading user turn for them too. Anthropic and Qwen # accept a leading assistant turn, so we leave their history untouched. # Note: several models reject whitespace-only text, so we use a period. if not processed_messages or ( processed_messages[0]['role'] != 'user' and not profile.get('bedrock_supports_leading_assistant_message', False) ): processed_messages.insert(0, {'role': 'user', 'content': [{'text': '.'}]}) return system_prompt, processed_messages @staticmethod def _get_last_user_message_content(messages: list[MessageUnionTypeDef]) -> list[Any] | None: """Get the content list from the last user message that can receive a cache point. Returns the content list if: - A user message exists - It has a non-empty content list - The last content block doesn't already have a cache point Returns None otherwise. """ user_messages = [msg for msg in messages if msg.get('role') == 'user'] if not user_messages: return None content = user_messages[-1].get('content') # Last user message if not content or not isinstance(content, list) or len(content) == 0: return None last_block = content[-1] if not isinstance(last_block, dict): return None if 'cachePoint' in last_block: # Skip if already has a cache point return None return content @staticmethod def _attach_cache_point_to_last_user_message( messages: list[MessageUnionTypeDef], cache_point: ContentBlockUnionTypeDef ) -> None: """Attach a cache point to the end of the last user message in `messages`. Used when a `CachePoint` is the first item of a part: everything the marker is meant to cache lives in the messages already built, so the marker goes there. If that message already ends with a cache point, the newest one replaces it, matching the Anthropic mapper. Raises: UserError: If there is no prior user message to attach the cache point to. """ content: list[Any] | None = next( ( message['content'] for message in reversed(messages) if message['role'] == 'user' and isinstance(message['content'], list) and message['content'] ), None, ) if content is None: raise UserError( 'CachePoint cannot be the first content in a user message - there must be previous content to cache when using Bedrock. ' 'To cache system instructions or tool definitions, use the `bedrock_cache_instructions` or `bedrock_cache_tool_definitions` settings instead.' ) last_block = content[-1] if isinstance(last_block, dict) and 'cachePoint' in last_block: content[-1] = cache_point else: _insert_cache_point_before_trailing_documents(content, cache_point, replace_existing=True) @staticmethod async def _map_file_to_content_block( file: ImageUrl | DocumentUrl | VideoUrl | BinaryContent, document_count: Iterator[int], ) -> ContentBlockUnionTypeDef: """Map a multimodal file directly to a Bedrock content block.""" source: DocumentSourceTypeDef if isinstance(file, BinaryContent): source = {'bytes': file.data} if file.is_image: return _make_image_block(file.format, source) elif file.is_document: return _make_document_block(f'Document {next(document_count)}', file.format, source) elif file.is_video: return _make_video_block(file.format, source) else: raise NotImplementedError(f'Unsupported binary content type for Bedrock: {file.media_type}') else: if file.url.startswith('s3://'): source = _parse_s3_source(file.url) else: downloaded = await download_item(file, data_format='bytes', type_format='extension') source = {'bytes': downloaded['data']} try: format = file.format except (KeyError, ValueError): format = file.media_type.split('/', 1)[1] if isinstance(file, ImageUrl): return _make_image_block(format, source) elif isinstance(file, DocumentUrl): return _make_document_block(f'Document {next(document_count)}', format, source) else: return _make_video_block(format, source) async def _map_user_prompt( # noqa: C901 self, part: UserPromptPart, document_count: Iterator[int], *, supports_prompt_caching: bool, prior_messages: list[MessageUnionTypeDef], ) -> list[MessageUnionTypeDef]: content: list[ContentBlockUnionTypeDef] = [] if isinstance(part.content, str): content.append({'text': part.content}) else: for item in part.content: if isinstance(item, str | TextContent): text = item if isinstance(item, str) else item.content content.append({'text': text}) elif isinstance(item, (BinaryContent, ImageUrl, DocumentUrl, VideoUrl)): content.append(await BedrockConverseModel._map_file_to_content_block(item, document_count)) elif isinstance(item, AudioUrl): raise NotImplementedError('AudioUrl is not supported in Bedrock user prompts') elif isinstance(item, UploadedFile): self._validate_uploaded_file_provider(item) if not item.file_id.startswith('s3://'): raise UserError( f'UploadedFile for Bedrock must use an S3 URL (s3://bucket/key), got: {item.file_id}' ) source: DocumentSourceTypeDef = _parse_s3_source(item.file_id) try: format = item.format except ValueError as e: raise UserError(f'Unsupported media type for Bedrock UploadedFile: {item.media_type}') from e if item.media_type.startswith('image/'): content.append(_make_image_block(format, source)) elif item.media_type.startswith('video/'): content.append(_make_video_block(format, source)) elif item.media_type.startswith('audio/'): raise UserError('Audio files are not supported for Bedrock UploadedFile') else: content.append(_make_document_block(f'Document {next(document_count)}', format, source)) elif isinstance(item, CachePoint): if not supports_prompt_caching: # Silently skip CachePoint for models that don't support prompt caching continue if not content: # A CachePoint means "cache everything before this point". This part has no # content yet, so everything before this point is the messages already built: # put the marker at the end of the last user message instead of raising. # See https://github.com/pydantic/pydantic-ai/issues/7004. self._attach_cache_point_to_last_user_message(prior_messages, self._get_cache_point(item.ttl)) continue if 'cachePoint' in content[-1]: raise UserError( 'CachePoint cannot be preceded by another CachePoint - there must be content between cache points when using Bedrock.' ) _insert_cache_point_before_trailing_documents( content, BedrockConverseModel._get_cache_point(item.ttl), raise_if_cannot_insert=True, ) else: assert_never(item) # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Message.html # "If you include a ContentBlock with a document field, you must also include a ContentBlock with a text field." has_document = any('document' in block for block in content) has_text = any('text' in block for block in content) if has_document and not has_text: content.insert(0, {'text': 'See attached document(s).'}) # A part that contained only a CachePoint has nothing left to send: emit no message. return [{'role': 'user', 'content': content}] if content else [] @staticmethod def _map_tool_call(t: ToolCallPart) -> ContentBlockOutputTypeDef: return { 'toolUse': { 'toolUseId': _utils.guard_tool_call_id(t=t), 'name': _utils.sanitize_tool_name(t.tool_name), 'input': t.args_as_dict(), } } @staticmethod def _get_cache_point(cache_setting: bool | Literal['5m', '1h']) -> ContentBlockUnionTypeDef: cache_point: CachePointBlockTypeDef = {'type': 'default'} if isinstance(cache_setting, str): cache_point['ttl'] = cache_setting return cast('ContentBlockUnionTypeDef', {'cachePoint': cache_point}) @staticmethod def _limit_cache_points( system_prompt: list[SystemContentBlockTypeDef], bedrock_messages: list[MessageUnionTypeDef], tools: list[ToolTypeDef], ) -> None: """Limit the number of cache points in the request to Bedrock's maximum. Bedrock enforces a maximum of 4 cache points per request. This method ensures compliance by counting existing cache points and removing excess ones from messages. Strategy: 1. Count cache points in system_prompt 2. Count cache points in tools 3. Raise UserError if system + tools already exceed MAX_CACHE_POINTS 4. Calculate remaining budget for message cache points 5. Traverse messages from newest to oldest, keeping the most recent cache points within the remaining budget 6. Remove excess cache points from older messages to stay within limit Cache point priority (always preserved): - System prompt cache points - Tool definition cache points - Message cache points (newest first, oldest removed if needed) Raises: UserError: If system_prompt and tools combined already exceed MAX_CACHE_POINTS (4). This indicates a configuration error that cannot be auto-fixed. """ MAX_CACHE_POINTS = 4 # Count existing cache points in system prompt used_cache_points = sum(1 for block in system_prompt if 'cachePoint' in block) # Count existing cache points in tools for tool in tools: if 'cachePoint' in tool: used_cache_points += 1 # Calculate remaining cache points budget for messages remaining_budget = MAX_CACHE_POINTS - used_cache_points if remaining_budget < 0: # pragma: no cover raise UserError( f'Too many cache points for Bedrock request. ' f'System prompt and tool definitions already use {used_cache_points} cache points, ' f'which exceeds the maximum of {MAX_CACHE_POINTS}.' ) # Remove excess cache points from messages (newest to oldest) for message in reversed(bedrock_messages): content = message.get('content') if not content or not isinstance(content, list): # pragma: no cover continue # Build a new content list, keeping only cache points within budget new_content: list[Any] = [] for block in reversed(content): # Process newest first is_cache_point = isinstance(block, dict) and 'cachePoint' in block if is_cache_point: if remaining_budget > 0: remaining_budget -= 1 new_content.append(block) else: new_content.append(block) message['content'] = list(reversed(new_content)) # Restore original order @dataclass class BedrockStreamedResponse(StreamedResponse): """Implementation of `StreamedResponse` for Bedrock models.""" _model_name: BedrockModelName _model_profile: BedrockModelProfile _event_stream: EventStream[ConverseStreamOutputTypeDef] _provider_name: str _model_id_namespace: str _provider_url: str _timestamp: datetime = field(default_factory=_utils.now_utc) _provider_response_id: str | None = None def get_stream_cancel_errors(self) -> tuple[type[BaseException], ...]: return (BotoCoreError, ClientError) async def close_stream(self) -> None: await anyio.to_thread.run_sync(self._event_stream.close) async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: # noqa: C901 with _map_api_errors(self._model_name, self._model_id_namespace): if self._provider_response_id is not None: self.provider_response_id = self._provider_response_id chunk: ConverseStreamOutputTypeDef tool_ids: dict[int, str] = {} # Bedrock has deltas for built-in tool returns, which aren't supported by parts manager. # We accumulate the deltas here and yield the complete return part once the content block ends builtin_tool_returns: dict[int, NativeToolReturnPart] = {} async for chunk in _AsyncIteratorWrapper(self._event_stream): match chunk: case {'messageStart': _}: continue case {'messageStop': message_stop}: raw_finish_reason = message_stop['stopReason'] self.provider_details = {**(self.provider_details or {}), 'finish_reason': raw_finish_reason} self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason) case {'metadata': metadata}: if 'usage' in metadata: # pragma: no branch self._usage += _map_usage( metadata['usage'], self._provider_name, self._provider_url, self._model_name ) if 'trace' in metadata: # `messageStop` usually precedes `metadata`, but AWS documents no strict event # ordering, so merge rather than overwrite the finish reason. self.provider_details = {**(self.provider_details or {}), 'trace': metadata['trace']} case {'contentBlockStart': content_block_start}: index = content_block_start['contentBlockIndex'] start = content_block_start['start'] if 'toolUse' in start: tool_use_start = start['toolUse'] tool_id = tool_use_start['toolUseId'] tool_ids[index] = tool_id tool_name = tool_use_start['name'] if tool_use_start.get('type') == 'server_tool_use': if tool_name == 'nova_code_interpreter': # pragma: no branch part = NativeToolCallPart( tool_name=CodeExecutionTool.kind, tool_call_id=tool_id, provider_name=self.provider_name, ) part.otel_metadata = {'code_arg_name': 'snippet', 'code_arg_language': 'python'} yield self._parts_manager.handle_part(vendor_part_id=index, part=part) elif maybe_event := self._parts_manager.handle_tool_call_delta( vendor_part_id=index, tool_name=tool_name, args=None, tool_call_id=tool_id, ): # pragma: no branch yield maybe_event elif 'toolResult' in start: # pragma: no branch tool_result_start = start['toolResult'] tool_id = tool_result_start['toolUseId'] if tool_result_start.get('type') == 'nova_code_interpreter_result': # pragma: no branch return_part = NativeToolReturnPart( provider_name=self.provider_name, tool_name=CodeExecutionTool.kind, content=None, tool_call_id=tool_id, provider_details={'status': tool_result_start['status']} if 'status' in tool_result_start else {}, ) builtin_tool_returns[index] = return_part # Don't yield anything yet - we wait for content block end case {'contentBlockDelta': content_block_delta}: index = content_block_delta['contentBlockIndex'] delta = content_block_delta['delta'] if 'reasoningContent' in delta: if redacted_content := delta['reasoningContent'].get('redactedContent'): for event in self._parts_manager.handle_thinking_delta( vendor_part_id=index, id='redacted_content', signature=redacted_content.decode('utf-8'), provider_name=self.provider_name, ): yield event else: signature = delta['reasoningContent'].get('signature') for event in self._parts_manager.handle_thinking_delta( vendor_part_id=index, content=delta['reasoningContent'].get('text'), signature=signature, provider_name=self.provider_name if signature else None, ): yield event if text := delta.get('text'): for event in self._parts_manager.handle_text_delta( vendor_part_id=index, content=text, ignore_leading_whitespace=self._model_profile.get( 'ignore_streamed_leading_whitespace', False ), ): yield event if 'toolUse' in delta: tool_use = delta['toolUse'] maybe_event = self._parts_manager.handle_tool_call_delta( vendor_part_id=index, tool_name=tool_use.get('name'), args=tool_use.get('input'), tool_call_id=tool_ids[index], ) if maybe_event: # pragma: no branch yield maybe_event if 'toolResult' in delta: # pragma: no branch if ( return_part := builtin_tool_returns.get(index) ) and return_part.tool_name == CodeExecutionTool.kind: # pragma: no branch # For now, only process `contentBlockDelta.toolResult` for Code Exe tool. if tr_content := delta['toolResult']: # pragma: no branch # Goal here is to convert to object form. # This assumes the first item is the relevant one. return_part.content = tr_content[0].get('json') # Don't yield anything yet - we wait for content block end case {'contentBlockStop': content_block_stop}: index = content_block_stop['contentBlockIndex'] if return_part := builtin_tool_returns.get(index): # Emit the complete built-in tool return only once when the block closes. yield self._parts_manager.handle_part(vendor_part_id=index, part=return_part) tool_ids.pop(index, None) builtin_tool_returns.pop(index, None) case _: # pragma: no cover pass # pyright wants match statements to be exhaustive @property def model_name(self) -> str: """Get the model name of the response.""" return self._model_name @property def provider_name(self) -> str: """Get the provider name.""" return self._provider_name @property def provider_url(self) -> str: """Get the provider base URL.""" return self._provider_url @property def timestamp(self) -> datetime: return self._timestamp def _map_usage(usage_data: TokenUsageTypeDef, provider: str, provider_url: str, model: str) -> usage.RequestUsage: details: dict[str, int] = { k: v for k, v in usage_data.items() if k not in _BEDROCK_USAGE_FIELDS if isinstance(v, int) } return usage.RequestUsage.extract( dict(model=remove_bedrock_geo_prefix(model), usage=usage_data), provider=provider, provider_url=provider_url, provider_fallback='bedrock', details=details or None, ) class _AsyncIteratorWrapper(Generic[T]): """Wrap a synchronous iterator in an async iterator.""" def __init__(self, sync_iterator: Iterable[T]): self.sync_iterator = iter(sync_iterator) def __aiter__(self): return self async def __anext__(self) -> T: try: return await anyio.to_thread.run_sync(next, self.sync_iterator) except RuntimeError as e: if type(e.__cause__) is StopIteration: raise StopAsyncIteration else: raise e # pragma: lax no cover def _effective_thinking_type( model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters | None, profile: BedrockModelProfile, ) -> Literal['adaptive', 'enabled'] | None: """Resolve the thinking type used by the output-tool and tool-forcing guards. Explicit `bedrock_additional_model_requests_fields` take precedence over unified thinking, matching `_build_additional_model_request_fields`. """ if model_settings and profile.get('bedrock_thinking_variant') == 'anthropic': additional_fields = model_settings.get('bedrock_additional_model_requests_fields') if additional_fields is not None and 'thinking' in additional_fields: thinking_config = additional_fields['thinking'] if _utils.is_str_dict(thinking_config): thinking_type = thinking_config.get('type') if thinking_type in ('adaptive', 'enabled'): return thinking_type return None if model_settings is not None and 'thinking' in model_settings: unified_thinking = model_settings['thinking'] else: unified_thinking = model_request_parameters.thinking if model_request_parameters is not None else None if not unified_thinking: return None if profile.get('bedrock_thinking_variant') == 'anthropic' and profile.get( 'bedrock_supports_adaptive_thinking', False ): return 'adaptive' return 'enabled' def _thinking_blocks_tool_forcing( thinking_type: Literal['adaptive', 'enabled'] | None, profile: BedrockModelProfile ) -> bool: return thinking_type == 'enabled' or (thinking_type == 'adaptive' and not _supports_tool_forcing(profile)) def _supports_tool_forcing(profile: BedrockModelProfile) -> bool: """Keep Anthropic's forcing capability distinct from Bedrock's general tool-choice support.""" supports_tool_choice = profile.get('bedrock_supports_tool_choice', False) if profile.get('bedrock_thinking_variant') == 'anthropic' and 'anthropic_supports_forced_tool_choice' in profile: return supports_tool_choice and bool(profile['anthropic_supports_forced_tool_choice']) return supports_tool_choice def _support_tool_forcing( model_name: str, profile: BedrockModelProfile, model_settings: BedrockModelSettings | None, model_request_parameters: ModelRequestParameters, effective_tool_choice: ResolvedToolChoice, ) -> bool: """Check if model supports tool forcing, raising UserError if explicitly requested but unsupported. Also checks thinking compatibility: extended thinking blocks forced tool choice, while adaptive thinking allows it on profiles that advertise support. """ if not _supports_tool_forcing(profile): explicit_choice = (model_settings or {}).get('tool_choice') if explicit_choice == 'required' or isinstance(explicit_choice, list): raise UserError( f'tool_choice={explicit_choice!r} is not supported by model {model_name!r}. ' f'This model does not support forcing tool use.' ) return False thinking_type = _effective_thinking_type(model_settings, model_request_parameters, profile) if _thinking_blocks_tool_forcing(thinking_type, profile): explicit_choice = (model_settings or {}).get('tool_choice') if explicit_choice == 'required' or isinstance(explicit_choice, list): context = "tool_choice='required'" if explicit_choice == 'required' else 'forcing specific tools' if profile.get('bedrock_thinking_variant') != 'anthropic': raise UserError( f'Bedrock does not support {context} with thinking enabled. ' f"Disable thinking or use `tool_choice='auto'`." ) adaptive_hint = ( f' Alternatively, `{_ADAPTIVE_THINKING_SETTING}` supports forcing.' if profile.get('bedrock_supports_adaptive_thinking', False) else '' ) raise UserError( f'Bedrock does not support {context} with extended thinking. ' f"Disable thinking or use `tool_choice='auto'`.{adaptive_hint}" ) if effective_tool_choice == 'required' or isinstance(effective_tool_choice, tuple): return False return True