diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py index 4141544098..4b70659622 100644 --- a/pydantic_ai_slim/pydantic_ai/models/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/models/bedrock.py @@ -72,6 +72,7 @@ PromptVariableValuesTypeDef, ReasoningContentBlockOutputTypeDef, S3LocationTypeDef, + ServiceTierTypeDef, SystemContentBlockTypeDef, ToolChoiceTypeDef, ToolConfigurationTypeDef, @@ -231,6 +232,12 @@ class BedrockModelSettings(ModelSettings, total=False): 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 + + See more about it on . + """ + @dataclass(init=False) class BedrockConverseModel(Model): @@ -481,6 +488,8 @@ async def _messages_create( params['additionalModelRequestFields'] = additional_model_requests_fields if prompt_variables := model_settings.get('bedrock_prompt_variables', None): params['promptVariables'] = prompt_variables + if service_tier := model_settings.get('bedrock_service_tier', None): + params['serviceTier'] = service_tier try: if stream: diff --git a/pydantic_ai_slim/pyproject.toml b/pydantic_ai_slim/pyproject.toml index 735c10d506..447e0f07a6 100644 --- a/pydantic_ai_slim/pyproject.toml +++ b/pydantic_ai_slim/pyproject.toml @@ -75,7 +75,7 @@ anthropic = ["anthropic>=0.75.0"] groq = ["groq>=0.25.0"] openrouter = ["openai>=2.8.0"] mistral = ["mistralai>=1.9.10"] -bedrock = ["boto3>=1.40.14"] +bedrock = ["boto3>=1.42.14"] huggingface = ["huggingface-hub[inference]>=0.33.5,<1.0.0"] outlines-transformers = ["outlines[transformers]>=1.0.0, <1.3.0; (sys_platform != 'darwin' or platform_machine != 'x86_64')", "transformers>=4.0.0", "pillow", "torch; (sys_platform != 'darwin' or platform_machine != 'x86_64')"] outlines-llamacpp = ["outlines[llamacpp]>=1.0.0, <1.3.0"] diff --git a/pyproject.toml b/pyproject.toml index ffd0050679..f7d0ae5ab4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,7 +104,7 @@ dev = [ "pytest-pretty>=1.3.0", "pytest-recording>=0.13.2", "diff-cover>=9.2.0", - "boto3-stubs[bedrock-runtime]", + "boto3-stubs[bedrock-runtime]>=1.42.13", "strict-no-cover @ git+https://github.com/pydantic/strict-no-cover.git@7fc59da2c4dff919db2095a0f0e47101b657131d", "pytest-xdist>=3.6.1", # Needed for PyCharm users diff --git a/tests/models/cassettes/test_bedrock/test_bedrock_model_service_tier.yaml b/tests/models/cassettes/test_bedrock/test_bedrock_model_service_tier.yaml new file mode 100644 index 0000000000..b6aa2ac12c --- /dev/null +++ b/tests/models/cassettes/test_bedrock/test_bedrock_model_service_tier.yaml @@ -0,0 +1,48 @@ +interactions: +- request: + body: '{"messages": [{"role": "user", "content": [{"text": "What is the capital of France?"}]}], "system": [{"text": "You + are a helpful chatbot."}], "serviceTier": {"type": "default"}}' + headers: + amz-sdk-invocation-id: + - !!binary | + ZmZhYjMyZmItODRjOS00YWZjLWE4NTAtNTQ4OTUxMjI5NmU4 + amz-sdk-request: + - !!binary | + YXR0ZW1wdD0x + content-length: + - '178' + content-type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.amazon.nova-micro-v1%3A0/converse + response: + headers: + connection: + - keep-alive + content-length: + - '206' + content-type: + - application/json + parsed_body: + metrics: + latencyMs: 159 + output: + message: + content: + - text: The capital of France is Paris. Paris is not only the capital city but also the most populous city in France, known for its significant cultural, + political, and economic influence both within the country and globally. It is famous for landmarks such as the Eiffel Tower, the Louvre Museum, + and the Notre-Dame Cathedral, among many other historical and architectural treasures. + + role: assistant + stopReason: max_tokens + usage: + inputTokens: 13 + outputTokens: 5 + totalTokens: 18 + serviceTier: + type: flex + status: + code: 200 + message: OK +version: 1 diff --git a/tests/models/test_bedrock.py b/tests/models/test_bedrock.py index 80e9d30d1d..23a4ffb527 100644 --- a/tests/models/test_bedrock.py +++ b/tests/models/test_bedrock.py @@ -562,6 +562,16 @@ async def test_bedrock_model_other_parameters(allow_model_requests: None, bedroc ) +async def test_bedrock_model_service_tier(allow_model_requests: None, bedrock_provider: BedrockProvider): + model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider) + model_settings = BedrockModelSettings(bedrock_service_tier={'type': 'flex'}) + agent = Agent(model=model, system_prompt='You are a helpful chatbot.', model_settings=model_settings) + result = await agent.run('What is the capital of France?') + assert result.output == snapshot( + 'The capital of France is Paris. Paris is not only the capital city but also the most populous city in France, known for its significant cultural, political, and economic influence both within the country and globally. It is famous for landmarks such as the Eiffel Tower, the Louvre Museum, and the Notre-Dame Cathedral, among many other historical and architectural treasures.' + ) + + async def test_bedrock_model_iter_stream(allow_model_requests: None, bedrock_provider: BedrockProvider): model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider) agent = Agent(model=model, system_prompt='You are a helpful chatbot.', model_settings={'top_p': 0.5}) diff --git a/uv.lock b/uv.lock index a1fc62b904..645affbf66 100644 --- a/uv.lock +++ b/uv.lock @@ -5964,7 +5964,7 @@ provides-extras = ["a2a", "dbos", "examples", "outlines-llamacpp", "outlines-mlx dev = [ { name = "anyio", specifier = ">=4.5.0" }, { name = "asgi-lifespan", specifier = ">=2.1.0" }, - { name = "boto3-stubs", extras = ["bedrock-runtime"] }, + { name = "boto3-stubs", extras = ["bedrock-runtime"], specifier = ">=1.42.13" }, { name = "coverage", extras = ["toml"], specifier = ">=7.10.7" }, { name = "devtools", specifier = ">=0.12.2" }, { name = "diff-cover", specifier = ">=9.2.0" }, @@ -6168,7 +6168,7 @@ requires-dist = [ { name = "ag-ui-protocol", marker = "extra == 'ag-ui'", specifier = ">=0.1.10" }, { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.75.0" }, { name = "argcomplete", marker = "extra == 'cli'", specifier = ">=3.5.0" }, - { name = "boto3", marker = "extra == 'bedrock'", specifier = ">=1.40.14" }, + { name = "boto3", marker = "extra == 'bedrock'", specifier = ">=1.42.14" }, { name = "cohere", marker = "sys_platform != 'emscripten' and extra == 'cohere'", specifier = ">=5.18.0" }, { name = "dbos", marker = "extra == 'dbos'", specifier = ">=1.14.0" }, { name = "ddgs", marker = "extra == 'duckduckgo'", specifier = ">=9.0.0" },