Skip to content
15 changes: 4 additions & 11 deletions pydantic_ai_slim/pydantic_ai/models/bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError, UserError
from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, download_item
from pydantic_ai.providers import Provider, infer_provider
from pydantic_ai.providers.bedrock import BedrockModelProfile
from pydantic_ai.providers.bedrock import BEDROCK_GEO_PREFIXES, BedrockModelProfile
from pydantic_ai.settings import ModelSettings
from pydantic_ai.tools import ToolDefinition

Expand Down Expand Up @@ -155,13 +155,6 @@
'tool_use': 'tool_call',
}

_AWS_BEDROCK_INFERENCE_GEO_PREFIXES: tuple[str, ...] = ('us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.', 'global.')
"""Geo prefixes for Bedrock inference profile IDs (e.g., 'eu.', 'us.').

Used to strip the geo prefix so we can pass a pure foundation model ID/ARN to CountTokens,
which does not accept profile IDs. Extend if new geos appear (e.g., 'global.', 'us-gov.').
"""


class BedrockModelSettings(ModelSettings, total=False):
"""Settings for Bedrock models.
Expand Down Expand Up @@ -693,9 +686,9 @@ def _map_tool_call(t: ToolCallPart) -> ContentBlockOutputTypeDef:
@staticmethod
def _remove_inference_geo_prefix(model_name: BedrockModelName) -> BedrockModelName:
"""Remove inference geographic prefix from model ID if present."""
for prefix in _AWS_BEDROCK_INFERENCE_GEO_PREFIXES:
if model_name.startswith(prefix):
return model_name.removeprefix(prefix)
for prefix in BEDROCK_GEO_PREFIXES:
if model_name.startswith(f'{prefix}.'):
return model_name.removeprefix(f'{prefix}.')
return model_name


Expand Down
9 changes: 7 additions & 2 deletions pydantic_ai_slim/pydantic_ai/providers/bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,10 @@ def bedrock_deepseek_model_profile(model_name: str) -> ModelProfile | None:
return profile # pragma: no cover


# Known geo prefixes for cross-region inference profile IDs
BEDROCK_GEO_PREFIXES: tuple[str, ...] = ('us', 'eu', 'apac', 'jp', 'au', 'ca', 'global', 'us-gov')


class BedrockProvider(Provider[BaseClient]):
"""Provider for AWS Bedrock."""

Expand Down Expand Up @@ -90,10 +94,11 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
# Split the model name into parts
parts = model_name.split('.', 2)

# Handle regional prefixes (e.g. "us.")
if len(parts) > 2 and len(parts[0]) == 2:
# Handle regional prefixes
if len(parts) > 2 and parts[0] in BEDROCK_GEO_PREFIXES:
parts = parts[1:]

# required format is provider.model-name-with-version
if len(parts) < 2:
return None

Expand Down
53 changes: 51 additions & 2 deletions tests/providers/test_bedrock.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import cast
from typing import cast, get_args

import pytest
from pytest_mock import MockerFixture
Expand All @@ -16,7 +16,8 @@
with try_import() as imports_successful:
from mypy_boto3_bedrock_runtime import BedrockRuntimeClient

from pydantic_ai.providers.bedrock import BedrockModelProfile, BedrockProvider
from pydantic_ai.models.bedrock import LatestBedrockModelNames
from pydantic_ai.providers.bedrock import BEDROCK_GEO_PREFIXES, BedrockModelProfile, BedrockProvider


pytestmark = pytest.mark.skipif(not imports_successful(), reason='bedrock not installed')
Expand Down Expand Up @@ -100,3 +101,51 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture):

unknown_model = provider.model_profile('unknown.unknown-model')
assert unknown_model is None


@pytest.mark.parametrize('prefix', BEDROCK_GEO_PREFIXES)
def test_bedrock_provider_model_profile_all_geo_prefixes(env: TestEnv, prefix: str):
"""Test that all cross-region inference geo prefixes are correctly handled."""
env.set('AWS_DEFAULT_REGION', 'us-east-1')
provider = BedrockProvider()

model_name = f'{prefix}.anthropic.claude-sonnet-4-5-20250929-v1:0'
profile = provider.model_profile(model_name)

assert profile is not None, f'model_profile returned None for {model_name}'


def test_bedrock_provider_model_profile_with_unknown_geo_prefix(env: TestEnv):
env.set('AWS_DEFAULT_REGION', 'us-east-1')
provider = BedrockProvider()

model_name = 'narnia.anthropic.claude-sonnet-4-5-20250929-v1:0'
profile = provider.model_profile(model_name)
assert profile is None, f'model_profile returned {profile} for {model_name}'
Comment on lines +118 to +124
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@DouweM this is in line with previous behavior but probably bad UX, right? should we parse the anthropic provider properly even if we haven't registered a potential new region prefix? and print a warning?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i don't know if we should properly parse if it's an incorrect geo prefix, as a user id be confused that a non existent geo prefix actually worked.

might be worth printing a warning and returning None still? or just raising an exception (unsure if that's a huge nono)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this behavior is fine, you'd end up getting an error from the API anyway



def test_latest_bedrock_model_names_geo_prefixes_are_supported():
"""Ensure all geo prefixes used in LatestBedrockModelNames are in BEDROCK_GEO_PREFIXES.

This test prevents adding new model names with geo prefixes that aren't handled
by the provider's model_profile method.
"""
model_names = get_args(LatestBedrockModelNames)

missing_prefixes: set[str] = set()

for model_name in model_names:
# Model names with geo prefixes have 3+ dot-separated parts:
# - No prefix: "anthropic.claude-xxx" (2 parts)
# - With prefix: "us.anthropic.claude-xxx" (3 parts)
parts = model_name.split('.')
if len(parts) >= 3:
geo_prefix = parts[0]
if geo_prefix not in BEDROCK_GEO_PREFIXES: # pragma: no cover
missing_prefixes.add(geo_prefix)

if missing_prefixes: # pragma: no cover
pytest.fail(
f'Found geo prefixes in LatestBedrockModelNames that are not in BEDROCK_GEO_PREFIXES: {missing_prefixes}. '
f'Please add them to BEDROCK_GEO_PREFIXES'
)