Repository: pydantic/pydantic-ai
File: tests/models/test_bedrock.py (Python)

Code before:
                        timestamp=IsDatetime(),
                    ),
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    TextPart(
                        content="Hello! How can I assist you today? Whether you have questions, need information, or just want to chat, I'm here to help."
                    )
                ],
                usage=RequestUsage(input_tokens=7, output_tokens=30),
                model_name='us.amazon.nova-micro-v1:0',
                timestamp=IsDatetime(),
                provider_name='bedrock',
                provider_url='https://bedrock-runtime.us-east-1.amazonaws.com',
                provider_details={'finish_reason': 'end_turn'},
                finish_reason='stop',
                run_id=IsStr(),
            ),
        ]
    )


@pytest.mark.parametrize(
    ('usage_payload', 'expected_usage'),
    [
        pytest.param(
            {
                'inputTokens': 3,
                'outputTokens': 5,
                'cacheReadInputTokens': 7,
                'cacheWriteInputTokens': 11,
            },
            RequestUsage(
                input_tokens=21,
                output_tokens=5,
                cache_read_tokens=7,
                cache_write_tokens=11,
            ),
            id='full-usage',
        ),
        pytest.param(
            {
                'outputTokens': 5,
                'cacheReadInputTokens': 7,
                'cacheWriteInputTokens': 11,
            },
            RequestUsage(
                input_tokens=18,

Review comment (suggestion):
these two tests are a bit too verbose and also rely on mocked behavior, I'd much prefer you rewriting existing cassettes and using them to assert the usage calculation or having a new test with fresh cassettes

Code after:
                usage=RequestUsage(input_tokens=7, output_tokens=30),
                model_name='us.amazon.nova-micro-v1:0',
                timestamp=IsDatetime(),
                provider_name='bedrock',
                provider_url='https://bedrock-runtime.us-east-1.amazonaws.com',
                provider_details={'finish_reason': 'end_turn'},
                finish_reason='stop',
                run_id=IsStr(),
            ),
        ]
    )


@pytest.mark.vcr()
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=23\\)',
    ):
        await agent.run(
            ['The quick brown fox jumps over the lazydog.', CachePoint(), 'What was next?'],
            usage_limits=UsageLimits(input_tokens_limit=18, count_tokens_before_request=True),
        )


@pytest.mark.vcr()
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?'