fix(anthropic): support Bedrock token counting by dyscott · Pull Request #5480 · pydantic/pydantic-ai

dyscott · 2026-05-15T21:36:30Z

Summary

Fix token counting for Anthropic models on Bedrock.
Remove the current AsyncAnthropicBedrock guard now that Bedrock supports CountTokens for Claude models.
Route Bedrock token counting through the supplied AsyncAnthropicBedrock client.
Build the existing Anthropic Messages request body, wrap it in Bedrock Runtime’s input.invokeModel.body CountTokens shape, base64-encode the inner body for the raw HTTP JSON protocol, and return BetaMessageTokensCount(input_tokens=...).
Caveat: I could not find another model implementation in this codebase that calls through to an underlying SDK client’s low- level post(...) method directly. I am not sure if what I implemented would be the best pattern going forward, so looking to hear feedback on this.
Closes Support token counting for Anthropic Bedrock models #5377

Testing Plan

New tests in tests/models/test_anthropic.py cover the Bedrock CountTokens request shape, including the base64-encoded
input.invokeModel.body, system, anthropic_betas, extra_body, and timeout handling.
Live-tested against anthropic.claude-haiku-4-5-20251001-v1:0 on Bedrock and got a successful token count response.
uv run pytest tests/models/test_anthropic.py::test_anthropic_bedrock_count_tokens tests/models/ test_anthropic.py::test_anthropic_bedrock_count_tokens_unexpected_response tests/models/ test_anthropic.py::test_anthropic_process_response_server_tool_blocks tests/models/ test_anthropic.py::test_anthropic_count_tokens_with_no_messages
uv run ruff check pydantic_ai_slim/pydantic_ai/models/anthropic.py tests/models/test_anthropic.py
uv run ruff format --check pydantic_ai_slim/pydantic_ai/models/anthropic.py tests/models/test_anthropic.py
uv run pyright pydantic_ai_slim/pydantic_ai/models/anthropic.py tests/models/test_anthropic.py

Checklist

Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it.
No breaking changes in accordance with the version policy.
PR title is fit for the release changelog.

Move the low-level Bedrock CountTokens logic out of AnthropicModel into models/_anthropic_bedrock_count_tokens.py, document why the SDK's high-level method can't be used, and replace the mock-only test with a recorded VCR test against real Bedrock (claude-sonnet-4 base model id).

…pic-bedrock-count-tokens # Conflicts: # pydantic_ai_slim/pydantic_ai/models/anthropic.py

dsfaccini · 2026-06-23T21:08:41Z

David's AICA here: thanks @dyscott — your approach is correct, and we verified it works against real Bedrock: count_tokens() returns input_tokens=18 via the low-level /model/{model}/count-tokens endpoint (the SDK blocks the high-level method on Bedrock, so this is the only path).

Cleanup we layered on top of your commits (authorship preserved):

Extracted the low-level Bedrock logic into its own module, pydantic_ai/models/_anthropic_bedrock_count_tokens.py, so it doesn't clutter AnthropicModel. There's a module comment documenting why the SDK's high-level method can't be used, so a future reader doesn't "clean it away".
Replaced the mock-only test with a recorded VCR test that exercises a real Bedrock round-trip (request envelope + parsed token count). The cassette uses anthropic.claude-sonnet-4-20250514-v1:0 rather than the claude-3-5-sonnet id from the original test, because that id is now end-of-life on Bedrock's CountTokens endpoint. Worth noting: CRIS inference-profile ids (us./eu./global. prefixes) aren't accepted there either — only base foundation-model ids — so we added a short docs note about that.
Merged latest main to resolve conflicts (your branch predated main's server-tools workaround for count_tokens, which we kept).

The approach itself is unchanged — purely a tidy-up plus the recorded test. Thanks again for the contribution!

coderabbitai · 2026-06-23T21:12:23Z

📝 Walkthrough

Walkthrough

The PR adds a Bedrock-specific token counting helper that posts to /model/{model}/count-tokens, validates inputTokens, and is used by AnthropicModel._messages_count_tokens for AsyncAnthropicBedrock. It also updates Bedrock token-count tests, removes the old unsupported-test assertion, adds recorded cassettes for the real API and error cases, and adds a documentation note about Bedrock model-id constraints.

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: Bedrock token counting support for Anthropic.
Description check	✅ Passed	The description matches the template well, including the issue link, summary, testing plan, and checklist.
Linked Issues check	✅ Passed	The implementation matches `#5377` by removing the Bedrock guard and routing token counting through Bedrock's count-tokens endpoint.
Out of Scope Changes check	✅ Passed	The added docs, helper extraction, cassettes, and tests are all directly related to the Bedrock token-counting fix.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/models/test_anthropic.py`:
- Around line 10451-10455: The AsyncAnthropicBedrock instantiation contains a
hardcoded secret-like literal for the aws_secret_key parameter, which triggers
Ruff S106. To fix this, either replace the hardcoded 'test-secret-key' value
with an environment variable lookup (using a pattern similar to the VCR test in
the file), or add a scoped # noqa: S106 comment on the line if the hardcoded
test value is intentional. Choose the approach that aligns with how other test
cases in this file handle sensitive test credentials.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: afe4ee65-5654-4ad6-9ff4-34ac96d6f691

📥 Commits

Reviewing files that changed from the base of the PR and between 15690e8 and 9bab46c.

📒 Files selected for processing (5)

docs/models/anthropic.md
pydantic_ai_slim/pydantic_ai/models/_anthropic_bedrock_count_tokens.py
pydantic_ai_slim/pydantic_ai/models/anthropic.py
tests/models/cassettes/test_anthropic/test_anthropic_bedrock_count_tokens_real_api.yaml
tests/models/test_anthropic.py

…n botocore - Use cast_to=object instead of dict[str, object]: the anthropic SDK passes cast_to to issubclass(), which raises TypeError on a subscripted generic under Python 3.10; validate the raw dict explicitly instead. - Skip the real-API VCR test when botocore is absent (AsyncAnthropicBedrock's SigV4 signer imports it at request time; not in the default install).

- Add `test_anthropic_bedrock_count_tokens_error`: a CRIS inference-profile id now surfaces Bedrock's own 400 as `ModelHTTPError` (the headline behavior), recorded live against Bedrock CountTokens. - Add `test_anthropic_bedrock_count_tokens_before_request`: exercises the user-facing `agent.run(count_tokens_before_request=True)` path, reusing the existing `_real_api` cassette. - Document why `test_anthropic_bedrock_count_tokens_unexpected_response` mocks the response instead of using a cassette. - Drop `test_anthropic_process_response_server_tool_blocks`: out-of-scope, tested pre-existing untouched code via a private method.

dsfaccini · 2026-06-24T22:47:33Z

David's AICA here: 🔍 Local review suite (automated, pre-merge)

Wave	Status	Warnings (defer/skip)
1 · code-correctness	✅ closed	2
2 · tests/cassettes/docs	✅ closed	1
3 · cross-cutting/siblings	✅ closed	0

Verdict: ✅ all 3 waves passed, 0 blocking, 0 required
Reviewed at 2026-06-25 · diff 31bcad04

Two earlier informational findings have now been addressed:

Test placement — the Bedrock count_tokens tests were extracted from the test_anthropic.py monolith into a dedicated tests/models/test_anthropic_bedrock_count_tokens.py feature file (cassettes moved alongside).
Server-tool counting (verified, not assumed) — probed live: Bedrock's InvokeModel /count-tokens endpoint rejects server tools (web_search) with HTTP 400, exactly like the standard endpoint, so the strip-to-MemoryTool is correct on Bedrock and there is no Bedrock-specific undercounting. Locked in by test_anthropic_bedrock_count_tokens_rejects_server_tools with a recorded 400.

⚠️ 2 informational (skip)

Automated pre-merge review — findings are pre-triage and may include false positives.

1. test_anthropic_bedrock_count_tokens_before_request reuses the _real_api cassette by filename · tests/models/test_anthropic_bedrock_count_tokens.py
@pytest.mark.vcr('...real_api.yaml') shares one deterministic count-tokens interaction; input_tokens_limit=10 raises UsageLimitExceeded before any message request, so only the count call replays — skip: reuse is correct and documented, avoids a redundant recording.

2. Body-key list can drift from the SDK count_tokens() kwargs · models/_anthropic_bedrock_count_tokens.py · count_tokens_via_bedrock
The Bedrock path hand-marshals the body from an explicit key list while the non-Bedrock path passes kwargs to the SDK; kept in sync only by the parallel call sites — skip: inherent to bypassing the SDK's Bedrock-blocked high-level method, no shared abstraction to factor out.

…ool rejection - Move the four Bedrock count_tokens tests (and their cassettes) out of the test_anthropic.py monolith into a dedicated test_anthropic_bedrock_count_tokens.py feature file. - Add `test_anthropic_bedrock_count_tokens_rejects_server_tools`: confirms (against a real recorded Bedrock response) that the InvokeModel /count-tokens endpoint 400s on a `web_search` server tool, exactly like the standard endpoint — so the strip-to-MemoryTool is necessary on Bedrock too and there is no Bedrock-specific undercounting.

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/models/test_anthropic_bedrock_count_tokens.py`:
- Around line 59-63: The test setup in AsyncAnthropicBedrock uses a literal
secret value that triggers Ruff S106, so update the aws_secret_key assignment in
the test to avoid a hardcoded secret. Either add the inline noqa suppression on
the aws_secret_key line or source the key from os.environ, keeping the change
localized to the AsyncAnthropicBedrock instantiation in the test.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: bdfa7283-a3cc-43d3-be99-65b4fa871d2b

📥 Commits

Reviewing files that changed from the base of the PR and between bd2f8d6 and 698f0d6.

📒 Files selected for processing (5)

tests/models/cassettes/test_anthropic_bedrock_count_tokens/test_anthropic_bedrock_count_tokens_error.yaml
tests/models/cassettes/test_anthropic_bedrock_count_tokens/test_anthropic_bedrock_count_tokens_real_api.yaml
tests/models/cassettes/test_anthropic_bedrock_count_tokens/test_anthropic_bedrock_count_tokens_rejects_server_tools.yaml
tests/models/test_anthropic.py
tests/models/test_anthropic_bedrock_count_tokens.py

💤 Files with no reviewable changes (3)

tests/models/cassettes/test_anthropic_bedrock_count_tokens/test_anthropic_bedrock_count_tokens_real_api.yaml
tests/models/cassettes/test_anthropic_bedrock_count_tokens/test_anthropic_bedrock_count_tokens_error.yaml
tests/models/test_anthropic.py

✅ Files skipped from review due to trivial changes (1)

tests/models/cassettes/test_anthropic_bedrock_count_tokens/test_anthropic_bedrock_count_tokens_rejects_server_tools.yaml

…tokens tests

…ts into try_import gate

Co-authored-by: Dylan Scott <dylan@sweetspot.so> Co-authored-by: David Sanchez <64162682+dsfaccini@users.noreply.github.com>

fix(anthropic): support Bedrock token counting

b33625c

github-actions Bot added size: M Medium PR (101-500 weighted lines) bug Report that something isn't working, or PR implementing a fix labels May 15, 2026

dyscott and others added 4 commits May 15, 2026 17:39

Merge branch 'main' into fix/anthropic-bedrock-count-tokens

2fe5b1f

fix(anthropic): support Bedrock token counting

dda8490

Merge remote-tracking branch 'upstream/main' into dyscott--fix/anthro…

9bab46c

…pic-bedrock-count-tokens # Conflicts: # pydantic_ai_slim/pydantic_ai/models/anthropic.py

coderabbitai Bot suggested changes Jun 23, 2026

View reviewed changes

Comment thread tests/models/test_anthropic.py Outdated

faccini-bs and others added 2 commits June 23, 2026 16:25

coderabbitai Bot approved these changes Jun 24, 2026

View reviewed changes

dsfaccini added the bedrock label Jun 25, 2026

coderabbitai Bot suggested changes Jun 25, 2026

View reviewed changes

Comment thread tests/models/test_anthropic_bedrock_count_tokens.py

dsfaccini added the douwebot reviews the PR that gets labeled and removes the label at the end label Jun 25, 2026

github-actions Bot reviewed Jun 25, 2026

View reviewed changes

Comment thread tests/models/test_anthropic_bedrock_count_tokens.py Outdated

github-actions Bot removed the douwebot reviews the PR that gets labeled and removes the label at the end label Jun 25, 2026

test(anthropic): extract bedrock_client fixture, dedup Bedrock count_…

ef0fbac

…tokens tests

dsfaccini reviewed Jun 25, 2026

View reviewed changes

Comment thread tests/models/test_anthropic_bedrock_count_tokens.py Outdated

dsfaccini reviewed Jun 25, 2026

View reviewed changes

Comment thread tests/models/test_anthropic_bedrock_count_tokens.py Outdated

test(anthropic): use single_request_body helper, hoist optional impor…

face6ad

…ts into try_import gate

dsfaccini merged commit 8375fd8 into pydantic:main Jun 25, 2026
64 checks passed

syf2211 added a commit to syf2211/pydantic-ai that referenced this pull request Jun 27, 2026

fix(anthropic): support Bedrock token counting (pydantic#5480)

9057b4b

Co-authored-by: Dylan Scott <dylan@sweetspot.so> Co-authored-by: David Sanchez <64162682+dsfaccini@users.noreply.github.com>

coderabbitai Bot mentioned this pull request Jul 7, 2026

Auto-enable Anthropic Files API beta when messages contain UploadedFile #5815

Merged

3 tasks

Navigation Menu

fix(anthropic): support Bedrock token counting - #5480

fix(anthropic): support Bedrock token counting#5480
dsfaccini merged 10 commits into
pydantic:mainfrom
dyscott:fix/anthropic-bedrock-count-tokens

dyscott commented May 15, 2026 •

edited

Loading

Uh oh!

dsfaccini commented Jun 23, 2026

Uh oh!

coderabbitai Bot commented Jun 23, 2026 •

edited

Loading

Walkthrough

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

dsfaccini commented Jun 24, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

dyscott commented May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing Plan

Checklist

Uh oh!

dsfaccini commented Jun 23, 2026

Uh oh!

coderabbitai Bot commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dsfaccini commented Jun 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dyscott commented May 15, 2026 •

edited

Loading

coderabbitai Bot commented Jun 23, 2026 •

edited

Loading

dsfaccini commented Jun 24, 2026 •

edited

Loading