From 879446f2f9aa2f19a89221baea549bed4d1c7904 Mon Sep 17 00:00:00 2001 From: Leundai Date: Thu, 4 Dec 2025 13:06:03 -0500 Subject: [PATCH 1/9] fix: support us-gov. and other multi-character Bedrock geo prefixes The BedrockProvider.model_profile method previously only handled 2-character regional prefixes (e.g., 'us.', 'eu.'), causing issues with models using longer prefixes like 'us-gov.' (AWS GovCloud) and 'global.'. This caused extended thinking to fail in multi-turn conversations because bedrock_send_back_thinking_parts stayed at its default False value for these models. ThinkingPart blocks from previous turns were converted to text blocks instead of reasoningContent, causing Bedrock to reject requests with: 'Expected thinking or redacted_thinking, but found text' Changes: - Add _strip_geo_prefix() helper function to properly handle all known geo prefixes including us-gov. and global. - Update _AWS_BEDROCK_INFERENCE_GEO_PREFIXES to include us-gov. - Add comprehensive tests for all geo prefixes Fixes cross-region inference for AWS GovCloud environments. --- .../pydantic_ai/models/bedrock.py | 15 ++++-- .../pydantic_ai/providers/bedrock.py | 27 ++++++++--- tests/providers/test_bedrock.py | 47 +++++++++++++++++++ 3 files changed, 80 insertions(+), 9 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py index 8d81948f68..c7528bf261 100644 --- a/pydantic_ai_slim/pydantic_ai/models/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/models/bedrock.py @@ -155,11 +155,20 @@ '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.'). +_AWS_BEDROCK_INFERENCE_GEO_PREFIXES: tuple[str, ...] = ( + 'us.', + 'eu.', + 'apac.', + 'jp.', + 'au.', + 'ca.', + 'global.', + 'us-gov.', +) +"""Geo prefixes for Bedrock inference profile IDs (e.g., 'eu.', 'us.', 'us-gov.'). 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.'). +which does not accept profile IDs. """ diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index f6fac74fae..2c48cfa1fe 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -58,6 +58,23 @@ 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.') + + +def _strip_geo_prefix(model_name: str) -> str: + """Strip geographic/regional prefix from model name if present. + + AWS Bedrock cross-region inference uses prefixes like 'us.', 'eu.', 'us-gov.', 'global.' + to route requests to specific regions. This function strips those prefixes so we can + identify the underlying provider and model. + """ + for prefix in _BEDROCK_GEO_PREFIXES: + if model_name.startswith(prefix): + return model_name.removeprefix(prefix) + return model_name + + class BedrockProvider(Provider[BaseClient]): """Provider for AWS Bedrock.""" @@ -87,13 +104,11 @@ def model_profile(self, model_name: str) -> ModelProfile | None: 'deepseek': bedrock_deepseek_model_profile, } - # 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: - parts = parts[1:] + # Strip regional/geo prefix if present (e.g. "us.", "eu.", "us-gov.", "global.") + model_name = _strip_geo_prefix(model_name) + # Split the model name into provider and model parts + parts = model_name.split('.', 1) if len(parts) < 2: return None diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index e791fe7a43..c249d5e134 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -100,3 +100,50 @@ 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', ['us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.', 'global.', 'us-gov.']) +def test_bedrock_provider_model_profile_all_geo_prefixes(env: TestEnv, prefix: str): + """Test that all cross-region inference geo prefixes are correctly handled. + + This is critical for AWS GovCloud (us-gov.) and other regional deployments + where models use prefixes longer than 2 characters. + """ + env.set('AWS_DEFAULT_REGION', 'us-east-1') + provider = BedrockProvider() + + # Test Anthropic model with geo prefix + 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}' + assert isinstance(profile, BedrockModelProfile) + assert profile.bedrock_supports_tool_choice is True + assert profile.bedrock_send_back_thinking_parts is True + + +def test_bedrock_provider_model_profile_us_gov_anthropic(env: TestEnv, mocker: MockerFixture): + """Test that us-gov. prefixed Anthropic models get the correct profile. + + This specifically tests the us-gov. prefix which was previously broken + because the provider only handled 2-character prefixes. + """ + env.set('AWS_DEFAULT_REGION', 'us-east-1') + provider = BedrockProvider() + + ns = 'pydantic_ai.providers.bedrock' + anthropic_model_profile_mock = mocker.patch(f'{ns}.anthropic_model_profile', wraps=anthropic_model_profile) + + # Test us-gov. prefix (AWS GovCloud cross-region inference) + profile = provider.model_profile('us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0') + anthropic_model_profile_mock.assert_called_with('claude-sonnet-4-5-20250929') + assert isinstance(profile, BedrockModelProfile) + assert profile.bedrock_supports_tool_choice is True + assert profile.bedrock_send_back_thinking_parts is True + + # Test global. prefix + profile = provider.model_profile('global.anthropic.claude-opus-4-5-20251101-v1:0') + anthropic_model_profile_mock.assert_called_with('claude-opus-4-5-20251101') + assert isinstance(profile, BedrockModelProfile) + assert profile.bedrock_supports_tool_choice is True + assert profile.bedrock_send_back_thinking_parts is True From 518dae7d0700ff4bd6e1d5975d7f2ad5cddbb611 Mon Sep 17 00:00:00 2001 From: Leundai Date: Sat, 6 Dec 2025 16:53:32 -0500 Subject: [PATCH 2/9] test: add validation that LatestBedrockModelNames geo prefixes are in BEDROCK_GEO_PREFIXES This test ensures we don't add new model names with geo prefixes that aren't handled by the provider's model_profile method. Also exports BEDROCK_GEO_PREFIXES in __all__. --- .../pydantic_ai/providers/bedrock.py | 7 ++- tests/providers/test_bedrock.py | 58 +++++++++---------- 2 files changed, 33 insertions(+), 32 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index 2c48cfa1fe..d66140739c 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -7,6 +7,8 @@ from typing import Any, Literal, overload from pydantic_ai import ModelProfile + +__all__ = ('BedrockModelProfile', 'BedrockProvider', 'BEDROCK_GEO_PREFIXES') from pydantic_ai.exceptions import UserError from pydantic_ai.profiles.amazon import amazon_model_profile from pydantic_ai.profiles.anthropic import anthropic_model_profile @@ -59,7 +61,7 @@ def bedrock_deepseek_model_profile(model_name: str) -> ModelProfile | None: # Known geo prefixes for cross-region inference profile IDs -_BEDROCK_GEO_PREFIXES: tuple[str, ...] = ('us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.', 'global.', 'us-gov.') +BEDROCK_GEO_PREFIXES: tuple[str, ...] = ('us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.', 'global.', 'us-gov.') def _strip_geo_prefix(model_name: str) -> str: @@ -69,7 +71,7 @@ def _strip_geo_prefix(model_name: str) -> str: to route requests to specific regions. This function strips those prefixes so we can identify the underlying provider and model. """ - for prefix in _BEDROCK_GEO_PREFIXES: + for prefix in BEDROCK_GEO_PREFIXES: if model_name.startswith(prefix): return model_name.removeprefix(prefix) return model_name @@ -104,7 +106,6 @@ def model_profile(self, model_name: str) -> ModelProfile | None: 'deepseek': bedrock_deepseek_model_profile, } - # Strip regional/geo prefix if present (e.g. "us.", "eu.", "us-gov.", "global.") model_name = _strip_geo_prefix(model_name) # Split the model name into provider and model parts diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index c249d5e134..59128a0b11 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -1,4 +1,4 @@ -from typing import cast +from typing import cast, get_args import pytest from pytest_mock import MockerFixture @@ -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') @@ -102,7 +103,7 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture): assert unknown_model is None -@pytest.mark.parametrize('prefix', ['us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.', 'global.', 'us-gov.']) +@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. @@ -117,33 +118,32 @@ def test_bedrock_provider_model_profile_all_geo_prefixes(env: TestEnv, prefix: s profile = provider.model_profile(model_name) assert profile is not None, f'model_profile returned None for {model_name}' - assert isinstance(profile, BedrockModelProfile) - assert profile.bedrock_supports_tool_choice is True - assert profile.bedrock_send_back_thinking_parts is True -def test_bedrock_provider_model_profile_us_gov_anthropic(env: TestEnv, mocker: MockerFixture): - """Test that us-gov. prefixed Anthropic models get the correct profile. +def test_latest_bedrock_model_names_geo_prefixes_are_supported(): + """Ensure all geo prefixes used in LatestBedrockModelNames are in BEDROCK_GEO_PREFIXES. - This specifically tests the us-gov. prefix which was previously broken - because the provider only handled 2-character prefixes. + This test prevents adding new model names with geo prefixes that aren't handled + by the provider's model_profile method. """ - env.set('AWS_DEFAULT_REGION', 'us-east-1') - provider = BedrockProvider() - - ns = 'pydantic_ai.providers.bedrock' - anthropic_model_profile_mock = mocker.patch(f'{ns}.anthropic_model_profile', wraps=anthropic_model_profile) - - # Test us-gov. prefix (AWS GovCloud cross-region inference) - profile = provider.model_profile('us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0') - anthropic_model_profile_mock.assert_called_with('claude-sonnet-4-5-20250929') - assert isinstance(profile, BedrockModelProfile) - assert profile.bedrock_supports_tool_choice is True - assert profile.bedrock_send_back_thinking_parts is True - - # Test global. prefix - profile = provider.model_profile('global.anthropic.claude-opus-4-5-20251101-v1:0') - anthropic_model_profile_mock.assert_called_with('claude-opus-4-5-20251101') - assert isinstance(profile, BedrockModelProfile) - assert profile.bedrock_supports_tool_choice is True - assert profile.bedrock_send_back_thinking_parts is True + model_names = get_args(LatestBedrockModelNames) + + # Known providers that appear after the geo prefix + known_providers = ('anthropic', 'amazon', 'meta', 'mistral', 'cohere', 'deepseek') + + missing_prefixes: set[str] = set() + + for model_name in model_names: + # Check if this model name has a geo prefix by seeing if it has 3+ dot-separated parts + # and the second part is a known provider + parts = model_name.split('.') + if len(parts) >= 3 and parts[1] in known_providers: + # This model has a geo prefix (e.g., "us.anthropic.claude...") + geo_prefix = parts[0] + '.' + if geo_prefix not in BEDROCK_GEO_PREFIXES: + missing_prefixes.add(geo_prefix) + + assert not missing_prefixes, ( + f'Found geo prefixes in LatestBedrockModelNames that are not in BEDROCK_GEO_PREFIXES: {missing_prefixes}. ' + f'Please add them to BEDROCK_GEO_PREFIXES in pydantic_ai/providers/bedrock.py' + ) From 745e2eb77a5451ec12a4b38b490bca38a547ef3f Mon Sep 17 00:00:00 2001 From: Leundai Date: Sat, 6 Dec 2025 16:57:32 -0500 Subject: [PATCH 3/9] refactor: simplify geo prefix detection without hardcoding providers Model names with geo prefixes have 3+ dot-separated parts, so we can detect them without knowing the provider names. --- tests/providers/test_bedrock.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index 59128a0b11..e552d7548a 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -128,17 +128,14 @@ def test_latest_bedrock_model_names_geo_prefixes_are_supported(): """ model_names = get_args(LatestBedrockModelNames) - # Known providers that appear after the geo prefix - known_providers = ('anthropic', 'amazon', 'meta', 'mistral', 'cohere', 'deepseek') - missing_prefixes: set[str] = set() for model_name in model_names: - # Check if this model name has a geo prefix by seeing if it has 3+ dot-separated parts - # and the second part is a known provider + # 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 and parts[1] in known_providers: - # This model has a geo prefix (e.g., "us.anthropic.claude...") + if len(parts) >= 3: geo_prefix = parts[0] + '.' if geo_prefix not in BEDROCK_GEO_PREFIXES: missing_prefixes.add(geo_prefix) From de437aca342e6881c3be4cbe45f914b0e995f4f4 Mon Sep 17 00:00:00 2001 From: Leundai Date: Sat, 6 Dec 2025 17:06:31 -0500 Subject: [PATCH 4/9] fix: again removing unused code --- .../pydantic_ai/models/bedrock.py | 20 ++----------------- 1 file changed, 2 insertions(+), 18 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py index c7528bf261..814fa235cc 100644 --- a/pydantic_ai_slim/pydantic_ai/models/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/models/bedrock.py @@ -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 @@ -155,22 +155,6 @@ 'tool_use': 'tool_call', } -_AWS_BEDROCK_INFERENCE_GEO_PREFIXES: tuple[str, ...] = ( - 'us.', - 'eu.', - 'apac.', - 'jp.', - 'au.', - 'ca.', - 'global.', - 'us-gov.', -) -"""Geo prefixes for Bedrock inference profile IDs (e.g., 'eu.', 'us.', 'us-gov.'). - -Used to strip the geo prefix so we can pass a pure foundation model ID/ARN to CountTokens, -which does not accept profile IDs. -""" - class BedrockModelSettings(ModelSettings, total=False): """Settings for Bedrock models. @@ -702,7 +686,7 @@ 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: + for prefix in BEDROCK_GEO_PREFIXES: if model_name.startswith(prefix): return model_name.removeprefix(prefix) return model_name From ee0ef61673aec4a3a2e0a6af07e4fc4bbcadff13 Mon Sep 17 00:00:00 2001 From: Leundai Date: Sat, 6 Dec 2025 17:09:17 -0500 Subject: [PATCH 5/9] fix: remove __all__ --- pydantic_ai_slim/pydantic_ai/providers/bedrock.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index d66140739c..909d55a84d 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -7,8 +7,6 @@ from typing import Any, Literal, overload from pydantic_ai import ModelProfile - -__all__ = ('BedrockModelProfile', 'BedrockProvider', 'BEDROCK_GEO_PREFIXES') from pydantic_ai.exceptions import UserError from pydantic_ai.profiles.amazon import amazon_model_profile from pydantic_ai.profiles.anthropic import anthropic_model_profile From ff6ebbd9e7aebbb708d32c2bbfb4b942938fbfbf Mon Sep 17 00:00:00 2001 From: Leundai Date: Sat, 6 Dec 2025 17:12:05 -0500 Subject: [PATCH 6/9] refactor: streamline docstring and error message in bedrock tests Updated the docstring for the test_bedrock_provider_model_profile_all_geo_prefixes function to be more concise. Modified the error message in the test_latest_bedrock_model_names_geo_prefixes_are_supported function to remove the file path reference, enhancing clarity. --- tests/providers/test_bedrock.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index e552d7548a..faabaf5c7d 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -105,15 +105,10 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture): @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. - - This is critical for AWS GovCloud (us-gov.) and other regional deployments - where models use prefixes longer than 2 characters. - """ + """Test that all cross-region inference geo prefixes are correctly handled.""" env.set('AWS_DEFAULT_REGION', 'us-east-1') provider = BedrockProvider() - # Test Anthropic model with geo prefix model_name = f'{prefix}anthropic.claude-sonnet-4-5-20250929-v1:0' profile = provider.model_profile(model_name) @@ -142,5 +137,5 @@ def test_latest_bedrock_model_names_geo_prefixes_are_supported(): assert not missing_prefixes, ( f'Found geo prefixes in LatestBedrockModelNames that are not in BEDROCK_GEO_PREFIXES: {missing_prefixes}. ' - f'Please add them to BEDROCK_GEO_PREFIXES in pydantic_ai/providers/bedrock.py' + f'Please add them to BEDROCK_GEO_PREFIXES' ) From 783ae2440170cb20c08adeafec12cafddd908f9f Mon Sep 17 00:00:00 2001 From: Leundai Date: Sat, 6 Dec 2025 17:15:36 -0500 Subject: [PATCH 7/9] fix: attempt to address coverage issue on test_bedrock --- tests/providers/test_bedrock.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index faabaf5c7d..bfb0020f3a 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -135,7 +135,8 @@ def test_latest_bedrock_model_names_geo_prefixes_are_supported(): if geo_prefix not in BEDROCK_GEO_PREFIXES: missing_prefixes.add(geo_prefix) - assert not missing_prefixes, ( - f'Found geo prefixes in LatestBedrockModelNames that are not in BEDROCK_GEO_PREFIXES: {missing_prefixes}. ' - f'Please add them to BEDROCK_GEO_PREFIXES' - ) + 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' + ) From b6067876ab5220e39cd6c9a22cd41b1070ec44b4 Mon Sep 17 00:00:00 2001 From: Leundai Date: Sat, 6 Dec 2025 17:43:29 -0500 Subject: [PATCH 8/9] fix: another attempt at the coverage issue here, --- tests/providers/test_bedrock.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index bfb0020f3a..667fcb6531 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -132,7 +132,7 @@ def test_latest_bedrock_model_names_geo_prefixes_are_supported(): parts = model_name.split('.') if len(parts) >= 3: geo_prefix = parts[0] + '.' - if geo_prefix not in BEDROCK_GEO_PREFIXES: + if geo_prefix not in BEDROCK_GEO_PREFIXES: # pragma: no cover missing_prefixes.add(geo_prefix) if missing_prefixes: # pragma: no cover From 8d8850bcb6af0d220c24c0b89eeb492cd96f677a Mon Sep 17 00:00:00 2001 From: Leundai Date: Mon, 8 Dec 2025 08:47:59 -0500 Subject: [PATCH 9/9] fix: simplify code change and constants This reduces the lines of code and also adjusts our tuple to now not have the period at the end. This will simplify our code. --- .../pydantic_ai/models/bedrock.py | 4 +-- .../pydantic_ai/providers/bedrock.py | 25 ++++++------------- tests/providers/test_bedrock.py | 13 ++++++++-- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py index 814fa235cc..57f517dad5 100644 --- a/pydantic_ai_slim/pydantic_ai/models/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/models/bedrock.py @@ -687,8 +687,8 @@ def _map_tool_call(t: ToolCallPart) -> ContentBlockOutputTypeDef: def _remove_inference_geo_prefix(model_name: BedrockModelName) -> BedrockModelName: """Remove inference geographic prefix from model ID if present.""" for prefix in BEDROCK_GEO_PREFIXES: - if model_name.startswith(prefix): - return model_name.removeprefix(prefix) + if model_name.startswith(f'{prefix}.'): + return model_name.removeprefix(f'{prefix}.') return model_name diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index 909d55a84d..3f855021e4 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -59,20 +59,7 @@ def bedrock_deepseek_model_profile(model_name: str) -> ModelProfile | None: # Known geo prefixes for cross-region inference profile IDs -BEDROCK_GEO_PREFIXES: tuple[str, ...] = ('us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.', 'global.', 'us-gov.') - - -def _strip_geo_prefix(model_name: str) -> str: - """Strip geographic/regional prefix from model name if present. - - AWS Bedrock cross-region inference uses prefixes like 'us.', 'eu.', 'us-gov.', 'global.' - to route requests to specific regions. This function strips those prefixes so we can - identify the underlying provider and model. - """ - for prefix in BEDROCK_GEO_PREFIXES: - if model_name.startswith(prefix): - return model_name.removeprefix(prefix) - return model_name +BEDROCK_GEO_PREFIXES: tuple[str, ...] = ('us', 'eu', 'apac', 'jp', 'au', 'ca', 'global', 'us-gov') class BedrockProvider(Provider[BaseClient]): @@ -104,10 +91,14 @@ def model_profile(self, model_name: str) -> ModelProfile | None: 'deepseek': bedrock_deepseek_model_profile, } - model_name = _strip_geo_prefix(model_name) + # Split the model name into parts + parts = model_name.split('.', 2) + + # Handle regional prefixes + if len(parts) > 2 and parts[0] in BEDROCK_GEO_PREFIXES: + parts = parts[1:] - # Split the model name into provider and model parts - parts = model_name.split('.', 1) + # required format is provider.model-name-with-version if len(parts) < 2: return None diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index 667fcb6531..c0dcbd2ddc 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -109,12 +109,21 @@ def test_bedrock_provider_model_profile_all_geo_prefixes(env: TestEnv, prefix: s env.set('AWS_DEFAULT_REGION', 'us-east-1') provider = BedrockProvider() - model_name = f'{prefix}anthropic.claude-sonnet-4-5-20250929-v1:0' + 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}' + + def test_latest_bedrock_model_names_geo_prefixes_are_supported(): """Ensure all geo prefixes used in LatestBedrockModelNames are in BEDROCK_GEO_PREFIXES. @@ -131,7 +140,7 @@ def test_latest_bedrock_model_names_geo_prefixes_are_supported(): # - With prefix: "us.anthropic.claude-xxx" (3 parts) parts = model_name.split('.') if len(parts) >= 3: - geo_prefix = parts[0] + '.' + geo_prefix = parts[0] if geo_prefix not in BEDROCK_GEO_PREFIXES: # pragma: no cover missing_prefixes.add(geo_prefix)