diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py index 81dc4f79638e..add4fb379476 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_model_info.py @@ -1,3 +1,4 @@ +import re from typing import Dict from autogen_core.models import ModelFamily, ModelInfo @@ -138,8 +139,26 @@ } +def _normalize_model(model: str) -> str: + """Normalize a Bedrock or cross-region inference model id to its bare Anthropic id. + + AWS Bedrock exposes Claude models with an ``anthropic.`` provider prefix, an optional + cross-region inference prefix (e.g. ``us.``, ``eu.``, ``apac.``, ``global.``), and a + ``-v:`` version suffix, e.g. ``us.anthropic.claude-sonnet-4-20250514-v1:0``. + Stripping these yields the bare id (``claude-sonnet-4-20250514``) so the lookup below can + resolve it the same way it resolves the native Anthropic id. Non-Bedrock ids are returned + unchanged. + """ + # Strip the optional cross-region prefix and the ``anthropic.`` provider prefix. + model = re.sub(r"^(?:[a-z0-9-]+\.)?anthropic\.", "", model) + # Strip the trailing Bedrock version suffix (e.g. ``-v1:0``). + model = re.sub(r"-v\d+:\d+$", "", model) + return model + + def get_info(model: str) -> ModelInfo: """Get the model information for a specific model.""" + model = _normalize_model(model) # Check for exact match first if model in _MODEL_INFO: return _MODEL_INFO[model] @@ -154,6 +173,7 @@ def get_info(model: str) -> ModelInfo: def get_token_limit(model: str) -> int: """Get the token limit for a specific model.""" + model = _normalize_model(model) # Check for exact match first if model in _MODEL_TOKEN_LIMITS: return _MODEL_TOKEN_LIMITS[model] diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py index 6238eee3bbe5..d11f9d0c882e 100644 --- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py @@ -41,6 +41,26 @@ def _ask_for_input() -> str: return "Further input from user" +@pytest.mark.parametrize( + "model", + [ + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "us.anthropic.claude-3-5-sonnet-20240620-v1:0", + "eu.anthropic.claude-3-5-sonnet-20240620-v1:0", + "apac.anthropic.claude-3-5-sonnet-20240620-v1:0", + "global.anthropic.claude-3-5-sonnet-20240620-v1:0", + ], +) +def test_model_info_resolves_bedrock_ids(model: str) -> None: + """Bedrock and cross-region inference ids resolve to the same entry as the bare Anthropic id.""" + from autogen_ext.models.anthropic import _model_info + + bare = "claude-3-5-sonnet-20240620" + assert _model_info.get_info(model) == _model_info.get_info(bare) + # The token limit is looked up from the table (200000), not the 100000 fallback. + assert _model_info.get_token_limit(model) == _model_info.get_token_limit(bare) == 200000 + + @pytest.mark.asyncio async def test_mock_tool_choice_specific_tool() -> None: """Test tool_choice parameter with a specific tool using mocks."""