Skip to content

feat(bedrock): support token counting#1650

Open
aidand-ant wants to merge 2 commits into
mainfrom
bedrock-count-tokens
Open

feat(bedrock): support token counting#1650
aidand-ant wants to merge 2 commits into
mainfrom
bedrock-count-tokens

Conversation

@aidand-ant
Copy link
Copy Markdown

Amazon Bedrock now supports token counting via the CountTokens API, so the SDK no longer needs to reject count_tokens() calls on the Bedrock client.

What changed

  • _prepare_options now rewrites POST /v1/messages/count_tokens (and the ?beta=true variant) to Bedrock's POST /model/{modelId}/count-tokens endpoint. The Anthropic-format request body is wrapped in Bedrock's CountTokensInput shape as a base64-encoded invokeModel body. The model ID is URL-quoted the same way as for create, so inference profile ARNs work.
  • BaseBedrockClient gains a _process_response_data override that maps Bedrock's {"inputTokens": N} response to {"input_tokens": N}, so callers get the same MessageTokensCount type as the first-party API.
  • The Bedrock beta messages resource previously only exposed create; it now also exposes count_tokens (including raw/streaming response wrappers).

Usage

from anthropic import AnthropicBedrock

client = AnthropicBedrock()
count = client.messages.count_tokens(
    model="anthropic.claude-3-5-sonnet-20241022-v2:0",
    messages=[{"role": "user", "content": "Hello, world"}],
)
print(count.input_tokens)

Tests

Added sync, async, and beta tests covering the URL rewrite, the base64-wrapped request body, and response parsing. All AWS-related tests pass (139 passed).

Amazon Bedrock now supports token counting via the CountTokens API,
so the Bedrock client no longer needs to reject count_tokens() calls.

- Rewrite POST /v1/messages/count_tokens (and the beta variant) to
  Bedrock's POST /model/{modelId}/count-tokens endpoint, wrapping the
  request body in the CountTokensInput shape as a base64-encoded
  invokeModel body
- Map Bedrock's inputTokens response field to input_tokens so callers
  get the same MessageTokensCount type as the first-party API
- Expose count_tokens on the Bedrock beta messages resource
mypy rejects wrapping a method assigned from another class with an
incompatible self type. Match the vertex pattern, which exposes
count_tokens on the resource but not on the raw/streaming response
wrappers.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant