Add BedrockConverseModel.count_tokens so it works with UsageLimits.count_tokens_before_request (pydantic#3367)

Co-authored-by: Douwe Maan <[email protected]>

Author: DenysMoskalenko

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -103,6 +103,13 @@
         'bedrock:us.anthropic.claude-opus-4-20250514-v1:0',
         'bedrock:anthropic.claude-sonnet-4-20250514-v1:0',
         'bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0',
+        'bedrock:eu.anthropic.claude-sonnet-4-20250514-v1:0',
+        'bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0',
+        'bedrock:us.anthropic.claude-sonnet-4-5-20250929-v1:0',
+        'bedrock:eu.anthropic.claude-sonnet-4-5-20250929-v1:0',
+        'bedrock:anthropic.claude-haiku-4-5-20251001-v1:0',
+        'bedrock:us.anthropic.claude-haiku-4-5-20251001-v1:0',
+        'bedrock:eu.anthropic.claude-haiku-4-5-20251001-v1:0',
         'bedrock:cohere.command-text-v14',
         'bedrock:cohere.command-r-v1:0',
         'bedrock:cohere.command-r-plus-v1:0',

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -22,7 +22,6 @@
     DocumentUrl,
     FinishReason,
     ImageUrl,
-    ModelHTTPError,
     ModelMessage,
     ModelProfileSpec,
     ModelRequest,
@@ -41,7 +40,7 @@
     usage,
 )
 from pydantic_ai._run_context import RunContext
-from pydantic_ai.exceptions import UserError
+from pydantic_ai.exceptions import 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
@@ -61,6 +60,7 @@
         ConverseStreamMetadataEventTypeDef,
         ConverseStreamOutputTypeDef,
         ConverseStreamResponseTypeDef,
+        CountTokensRequestTypeDef,
         DocumentBlockTypeDef,
         GuardrailConfigurationTypeDef,
         ImageBlockTypeDef,
@@ -77,7 +77,6 @@
         VideoBlockTypeDef,
     )
 
-
 LatestBedrockModelNames = Literal[
     'amazon.titan-tg1-large',
     'amazon.titan-text-lite-v1',
@@ -106,6 +105,13 @@
     'us.anthropic.claude-opus-4-20250514-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-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',
@@ -136,7 +142,6 @@
 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')
 
@@ -149,6 +154,13 @@
     'tool_use': 'tool_call',
 }
 
+_AWS_BEDROCK_INFERENCE_GEO_PREFIXES: tuple[str, ...] = ('us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.')
+"""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.
@@ -275,6 +287,34 @@ async def request(
         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 <https://docs.aws.amazon.com/bedrock/latest/userguide/count-tokens.html>
+        """
+        model_settings, model_request_parameters = self.prepare_request(model_settings, model_request_parameters)
+        system_prompt, bedrock_messages = await self._map_messages(messages, model_request_parameters)
+        params: CountTokensRequestTypeDef = {
+            'modelId': self._remove_inference_geo_prefix(self.model_name),
+            'input': {
+                'converse': {
+                    'messages': bedrock_messages,
+                    'system': system_prompt,
+                },
+            },
+        }
+        try:
+            response = await anyio.to_thread.run_sync(functools.partial(self.client.count_tokens, **params))
+        except ClientError as e:
+            status_code = e.response.get('ResponseMetadata', {}).get('HTTPStatusCode', 500)
+            raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.response) from e
+        return usage.RequestUsage(input_tokens=response['inputTokens'])
+
     @asynccontextmanager
     async def request_stream(
         self,
@@ -642,6 +682,14 @@ def _map_tool_call(t: ToolCallPart) -> ContentBlockOutputTypeDef:
             'toolUse': {'toolUseId': _utils.guard_tool_call_id(t=t), 'name': t.tool_name, 'input': t.args_as_dict()}
         }
 
+    @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)
+        return model_name
+
 
 @dataclass
 class BedrockStreamedResponse(StreamedResponse):

pydantic_ai_slim/pyproject.toml

@@ -74,7 +74,7 @@ google = ["google-genai>=1.46.0"]
 anthropic = ["anthropic>=0.70.0"]
 groq = ["groq>=0.25.0"]
 mistral = ["mistralai>=1.9.10"]
-bedrock = ["boto3>=1.39.0"]
+bedrock = ["boto3>=1.40.14"]
 huggingface = ["huggingface-hub[inference]>=0.33.5"]
 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"]

tests/conftest.py

@@ -424,6 +424,7 @@ def bedrock_provider():
             region_name=os.getenv('AWS_REGION', 'us-east-1'),
             aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID', 'AKIA6666666666666666'),
             aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY', '6666666666666666666666666666666666666666'),
+            aws_session_token=os.getenv('AWS_SESSION_TOKEN', None),
         )
         yield BedrockProvider(bedrock_client=bedrock_client)
         bedrock_client.close()

tests/models/cassettes/test_bedrock/test_bedrock_count_tokens_error.yaml

@@ -0,0 +1,31 @@
+interactions:
+- request:
+    body: '{"input": {"converse": {"messages": [{"role": "user", "content": [{"text": "hello"}]}], "system": []}}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        ODdjZWFjMTYtN2U4OC00YTMzLTg5Y2QtZDUwNWM4N2YzNmNk
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '103'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/does-not-exist-model-v1%3A0/count-tokens
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '55'
+      content-type:
+      - application/json
+    parsed_body:
+      message: The provided model identifier is invalid.
+    status:
+      code: 400
+      message: Bad Request
+version: 1

tests/models/cassettes/test_bedrock/test_bedrock_model_usage_limit_exceeded.yaml

@@ -0,0 +1,32 @@
+interactions:
+- request:
+    body: '{"input": {"converse": {"messages": [{"role": "user", "content": [{"text": "The quick brown fox jumps over the
+      lazydog."}]}], "system": []}}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        ZDYxNmVkOTktYzgwMi00MDE0LTljZGUtYWFjMjk5N2I2MDFj
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '141'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-20250514-v1%3A0/count-tokens
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '18'
+      content-type:
+      - application/json
+    parsed_body:
+      inputTokens: 19
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/cassettes/test_bedrock/test_bedrock_model_usage_limit_not_exceeded.yaml

@@ -0,0 +1,78 @@
+interactions:
+- request:
+    body: '{"input": {"converse": {"messages": [{"role": "user", "content": [{"text": "The quick brown fox jumps over the
+      lazydog."}]}], "system": []}}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        OWQ3NzFhZmItYTkwYi00N2E4LWFkNjMtZmI5OTJhZDEyN2E4
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '141'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-20250514-v1%3A0/count-tokens
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '18'
+      content-type:
+      - application/json
+    parsed_body:
+      inputTokens: 19
+    status:
+      code: 200
+      message: OK
+- request:
+    body: '{"messages": [{"role": "user", "content": [{"text": "The quick brown fox jumps over the lazydog."}]}], "system":
+      [], "inferenceConfig": {}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        MWMwNDdlYWEtOWIxMy00YjAyLWI3ZjMtMjZkNjQ2MDEzOTY2
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '139'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-sonnet-4-20250514-v1%3A0/converse
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '785'
+      content-type:
+      - application/json
+    parsed_body:
+      metrics:
+        latencyMs: 2333
+      output:
+        message:
+          content:
+          - text: "I notice there's a small typo in your message - it should be \"lazy dog\" (two words) rather than \"lazydog.\"\n\nThe corrected version is: \"The quick brown fox jumps over the lazy dog.\"\n\nThis is a famous pangram - a sentence that contains every letter of the English alphabet at least once. It's commonly used for testing typewriters, keyboards, fonts, and other applications where you want to display all the letters.\n\nIs there something specific you'd like to know about this phrase, or were you perhaps testing something?"
+          role: assistant
+      stopReason: end_turn
+      usage:
+        cacheReadInputTokenCount: 0
+        cacheReadInputTokens: 0
+        cacheWriteInputTokenCount: 0
+        cacheWriteInputTokens: 0
+        inputTokens: 19
+        outputTokens: 108
+        serverToolUsage: {}
+        totalTokens: 127
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/test_bedrock.py

@@ -14,7 +14,6 @@
     FunctionToolCallEvent,
     FunctionToolResultEvent,
     ImageUrl,
-    ModelHTTPError,
     ModelRequest,
     ModelResponse,
     PartDeltaEvent,
@@ -33,12 +32,12 @@
     VideoUrl,
 )
 from pydantic_ai.agent import Agent
-from pydantic_ai.exceptions import ModelRetry
+from pydantic_ai.exceptions import ModelHTTPError, ModelRetry, UsageLimitExceeded
 from pydantic_ai.messages import AgentStreamEvent
 from pydantic_ai.models import ModelRequestParameters
 from pydantic_ai.run import AgentRunResult, AgentRunResultEvent
 from pydantic_ai.tools import ToolDefinition
-from pydantic_ai.usage import RequestUsage, RunUsage
+from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
 
 from ..conftest import IsDatetime, IsInstance, IsStr, try_import
 
@@ -98,6 +97,73 @@ async def test_bedrock_model(allow_model_requests: None, bedrock_provider: Bedro
     )
 
 
+async def test_bedrock_model_usage_limit_exceeded(
+    allow_model_requests: None,
+    bedrock_provider: BedrockProvider,
+):
+    model = BedrockConverseModel('us.anthropic.claude-sonnet-4-20250514-v1:0', provider=bedrock_provider)
+    agent = Agent(model=model)
+
+    with pytest.raises(
+        UsageLimitExceeded,
+        match='The next request would exceed the input_tokens_limit of 18 \\(input_tokens=19\\)',
+    ):
+        await agent.run(
+            'The quick brown fox jumps over the lazydog.',
+            usage_limits=UsageLimits(input_tokens_limit=18, count_tokens_before_request=True),
+        )
+
+
+async def test_bedrock_model_usage_limit_not_exceeded(
+    allow_model_requests: None,
+    bedrock_provider: BedrockProvider,
+):
+    model = BedrockConverseModel('us.anthropic.claude-sonnet-4-20250514-v1:0', provider=bedrock_provider)
+    agent = Agent(model=model)
+
+    result = await agent.run(
+        'The quick brown fox jumps over the lazydog.',
+        usage_limits=UsageLimits(input_tokens_limit=25, count_tokens_before_request=True),
+    )
+
+    assert result.output == snapshot(
+        'I notice there\'s a small typo in your message - it should be "lazy dog" (two words) rather than '
+        '"lazydog."\n\nThe corrected version is: "The quick brown fox jumps over the lazy dog."\n\n'
+        'This is a famous pangram - a sentence that contains every letter of the English alphabet at least once. '
+        "It's commonly used for testing typewriters, keyboards, fonts, and other applications where you want to "
+        "display all the letters.\n\nIs there something specific you'd like to know about this phrase, or were you "
+        'perhaps testing something?'
+    )
+
+
+@pytest.mark.vcr()
+async def test_bedrock_count_tokens_error(allow_model_requests: None, bedrock_provider: BedrockProvider):
+    """Test that errors convert to ModelHTTPError."""
+    model_id = 'us.does-not-exist-model-v1:0'
+    model = BedrockConverseModel(model_id, provider=bedrock_provider)
+    agent = Agent(model)
+
+    with pytest.raises(ModelHTTPError) as exc_info:
+        await agent.run('hello', usage_limits=UsageLimits(input_tokens_limit=20, count_tokens_before_request=True))
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.model_name == model_id
+    assert exc_info.value.body.get('Error', {}).get('Message') == 'The provided model identifier is invalid.'  # type: ignore[union-attr]
+
+
+@pytest.mark.parametrize(
+    ('model_name', 'expected'),
+    [
+        ('us.anthropic.claude-sonnet-4-20250514-v1:0', 'anthropic.claude-sonnet-4-20250514-v1:0'),
+        ('eu.amazon.nova-micro-v1:0', 'amazon.nova-micro-v1:0'),
+        ('apac.meta.llama3-8b-instruct-v1:0', 'meta.llama3-8b-instruct-v1:0'),
+        ('anthropic.claude-3-7-sonnet-20250219-v1:0', 'anthropic.claude-3-7-sonnet-20250219-v1:0'),
+    ],
+)
+def test_remove_inference_geo_prefix(model_name: str, expected: str):
+    assert BedrockConverseModel._remove_inference_geo_prefix(model_name) == expected  # pyright: ignore[reportPrivateUsage]
+
+
 async def test_bedrock_model_structured_output(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.', retries=5)