Implement AWS Bedrock CountTokens API support

[timelfrink]

Title

Implement AWS Bedrock CountTokens API support

Relevant issues

Fixes #14478

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/litellm/ directory, Adding at least 1 test is a hard requirement - see details
• I have added a screenshot of my new test passing locally
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem

Type

✨ New Feature

Changes

Problem Description

Issue #14478 reported that the AWS Bedrock /v1/messages/count_tokens endpoint was not working properly. The original issue was that the endpoint path parsing was incorrectly extracting "messages" as the model name instead of properly handling the Anthropic-style API format.

Complete Solution Implementation

This PR provides a comprehensive implementation of AWS Bedrock CountTokens API support with proper endpoint handling, request transformation, and response formatting.

Key Components

1. Bedrock CountTokens Handler (litellm/llms/bedrock/count_tokens/handler.py)

• Implements the core token counting logic for Bedrock
• Supports both Converse API and InvokeModel formats
• Handles different input types (messages, raw text)
• Provides proper error handling and response formatting

2. Request Transformation (litellm/llms/bedrock/count_tokens/transformation.py)

• Transforms incoming requests to proper Bedrock format
• Handles model resolution and validation
• Supports both message-based and text-based inputs
• Maintains compatibility with existing LiteLLM patterns

3. Endpoint Integration (litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py)

• Adds /bedrock/v1/messages/count_tokens endpoint support
• Implements proper request routing and transformation
• Integrates with existing Bedrock authentication and model resolution
• Maintains backward compatibility with other Bedrock endpoints

Features

• Multiple Input Formats: Supports both Anthropic-style messages and raw text inputs
• Model Resolution: Automatic model name resolution using existing router patterns
• Authentication: Full integration with LiteLLM proxy authentication
• Error Handling: Comprehensive error handling with proper HTTP status codes
• Response Formatting: Returns token counts in expected format

API Usage

# Count tokens for messages
curl -X POST "http://localhost:4000/bedrock/v1/messages/count_tokens" \
  -H "Authorization: Bearer your-api-key" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "anthropic.claude-3-sonnet-20240229-v1:0",
    "messages": [
      {"role": "user", "content": "Hello, how are you?"}
    ]
  }'

# Response
{
  "input_tokens": 12
}

Implementation Details

Request Flow

1. Request comes to /bedrock/v1/messages/count_tokens
2. Endpoint is detected and routed to Bedrock count tokens handler
3. Request is transformed to proper Bedrock format
4. Model is resolved using existing router logic
5. Token counting is performed
6. Response is formatted and returned

Model Resolution

• Uses existing LiteLLM router patterns for model resolution
• Supports both full model names and aliases
• Maintains consistency with other LiteLLM endpoints

Error Handling

• Validates required fields (model, input content)
• Provides clear error messages for missing or invalid data
• Returns appropriate HTTP status codes
• Maintains consistency with LiteLLM error patterns

Testing

The implementation includes comprehensive testing to ensure reliability and compatibility.

Backward Compatibility

• No changes to existing Bedrock endpoints
• Maintains all existing authentication and routing logic
• Follows established LiteLLM patterns and conventions
• No breaking changes to existing functionality

Impact

• Enables: Full AWS Bedrock CountTokens API support through LiteLLM proxy
• Fixes: Issue [Bug]: bedrock passthrough endpoint returns error with claude code #14478 with proper endpoint handling
• Maintains: Complete backward compatibility
• Provides: Production-ready token counting for Bedrock models

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Preview	Comments	Updated (UTC)
litellm	Error			Sep 18, 2025 6:35am

[krrishdholakia]

@timelfrink have you manually qa'ed this pr? if so, can you please attach a video/screenshot of it working

[timelfrink]

@krrishdholakia updated the PR

Here validation:

Just checking, Vercel deployment failed, but I assume this is unrelated?

[krrishdholakia]

Thank you for manually qa'ing this @timelfrink