diff --git a/BEDROCK_SUPPORT.md b/BEDROCK_SUPPORT.md new file mode 100644 index 0000000..47e6752 --- /dev/null +++ b/BEDROCK_SUPPORT.md @@ -0,0 +1,83 @@ +# AWS Bedrock Model Support for tokencost + +## Problem Solved +Users integrating AgentOps with AWS Bedrock were unable to track costs for their LLM usage. The model identifier `bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0` was not recognized by the tokencost library, resulting in warnings like "Unable to calculate cost - This might be because you're using an unrecognized model." + +## Solution Implemented + +### 1. Added Model Name Normalization +Created a new function `normalize_bedrock_model_name()` in `/workspace/tokencost/costs.py` that: +- Strips the `bedrock/` prefix from model names +- Maps Bedrock model identifiers to their corresponding entries in the model prices dictionary + +### 2. Updated Cost Calculation Functions +Modified the following functions to use the normalization: +- `calculate_cost_by_tokens()` +- `calculate_prompt_cost()` +- `calculate_completion_cost()` +- `count_message_tokens()` +- `count_string_tokens()` + +### 3. Added Explicit Bedrock Model Entries +Added explicit entries in `/workspace/tokencost/model_prices.json` for: +- `bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0` +- `bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0` + +These entries have the correct pricing: +- Input tokens: $3.00 per 1M tokens (3e-06 per token) +- Output tokens: $15.00 per 1M tokens (1.5e-05 per token) +- Cached input tokens (v2 only): $0.30 per 1M tokens (3e-07 per token) + +### 4. Added Comprehensive Tests +Created test suite in `/workspace/tests/test_bedrock_models.py` that verifies: +- Cost calculation for Bedrock models with `bedrock/` prefix +- Cost calculation for models without the prefix +- Case-insensitive model name handling +- Cached token cost calculation +- Proper error handling for invalid models + +## Supported Model Formats + +The following AWS Bedrock model formats are now supported: + +### Claude 3.5 Sonnet Models +- `bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0` +- `bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0` +- `anthropic.claude-3-5-sonnet-20240620-v1:0` (without prefix) +- `anthropic.claude-3-5-sonnet-20241022-v2:0` (without prefix) + +### Other Bedrock Models +The normalization function will automatically handle any model with the `bedrock/` prefix by stripping it and looking up the base model name in the prices dictionary. + +## Usage Example + +```python +from tokencost import calculate_cost_by_tokens + +# Works with bedrock/ prefix +model = "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" +input_cost = calculate_cost_by_tokens(1000, model, "input") # Returns $0.003 +output_cost = calculate_cost_by_tokens(1000, model, "output") # Returns $0.015 + +# Also works without prefix +model = "anthropic.claude-3-5-sonnet-20240620-v1:0" +input_cost = calculate_cost_by_tokens(1000, model, "input") # Returns $0.003 +``` + +## Impact + +This fix enables: +- Cost tracking for AWS Bedrock users in AgentOps +- Budget management for production deployments using AWS Bedrock +- Support for CrewAI applications using AWS Bedrock LLM integration +- Compatibility with any framework using Bedrock model identifiers + +## Testing + +Run the test suite to verify the implementation: + +```bash +python3 -m pytest tests/test_bedrock_models.py -v +``` + +All tests should pass, confirming that AWS Bedrock models are now properly supported for cost calculation. \ No newline at end of file diff --git a/tests/test_bedrock_models.py b/tests/test_bedrock_models.py new file mode 100644 index 0000000..3e85e43 --- /dev/null +++ b/tests/test_bedrock_models.py @@ -0,0 +1,111 @@ +""" +Test AWS Bedrock model support in tokencost +""" + +import pytest +from decimal import Decimal +from tokencost import calculate_cost_by_tokens, calculate_prompt_cost, calculate_completion_cost + + +class TestBedrockModels: + """Test that AWS Bedrock model names are properly handled.""" + + def test_bedrock_claude_35_sonnet_v1_cost_calculation(self): + """Test cost calculation for bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0""" + model = "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" + + # Test input token cost: $3.00 per 1M tokens = 3e-06 per token + input_cost = calculate_cost_by_tokens(1000, model, "input") + assert input_cost == Decimal("0.003") + + # Test output token cost: $15.00 per 1M tokens = 1.5e-05 per token + output_cost = calculate_cost_by_tokens(1000, model, "output") + assert output_cost == Decimal("0.015") + + def test_bedrock_claude_35_sonnet_v2_cost_calculation(self): + """Test cost calculation for bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0""" + model = "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0" + + # Test input token cost + input_cost = calculate_cost_by_tokens(1000, model, "input") + assert input_cost == Decimal("0.003") + + # Test output token cost + output_cost = calculate_cost_by_tokens(1000, model, "output") + assert output_cost == Decimal("0.015") + + def test_bedrock_model_without_prefix(self): + """Test that models without bedrock/ prefix still work""" + model = "anthropic.claude-3-5-sonnet-20240620-v1:0" + + # Should work the same way + input_cost = calculate_cost_by_tokens(1000, model, "input") + assert input_cost == Decimal("0.003") + + output_cost = calculate_cost_by_tokens(1000, model, "output") + assert output_cost == Decimal("0.015") + + def test_bedrock_model_case_insensitive(self): + """Test that model names are case-insensitive""" + model_upper = "BEDROCK/ANTHROPIC.CLAUDE-3-5-SONNET-20240620-V1:0" + model_mixed = "Bedrock/Anthropic.Claude-3-5-Sonnet-20240620-v1:0" + + # Both should work + input_cost_upper = calculate_cost_by_tokens(1000, model_upper, "input") + input_cost_mixed = calculate_cost_by_tokens(1000, model_mixed, "input") + + assert input_cost_upper == Decimal("0.003") + assert input_cost_mixed == Decimal("0.003") + + def test_bedrock_prompt_cost_calculation(self): + """Test calculate_prompt_cost with Bedrock model""" + model = "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" + + # Note: For Claude models, this will require anthropic API key + # So we test with a simple prompt that would work with tiktoken fallback + # The actual token counting for Claude requires the anthropic library + try: + # This might fail without proper Anthropic API key + prompt = "Hello, world!" + cost = calculate_prompt_cost(prompt, model) + # Just verify it returns a Decimal, actual value depends on tokenization + assert isinstance(cost, Decimal) + except Exception as e: + # Expected if Anthropic API key is not set + if "ANTHROPIC_API_KEY" in str(e) or "API" in str(e): + pytest.skip("Anthropic API key not available for testing") + else: + raise + + def test_bedrock_completion_cost_calculation(self): + """Test calculate_completion_cost with Bedrock model""" + model = "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" + + try: + completion = "This is a test completion." + cost = calculate_completion_cost(completion, model) + # Just verify it returns a Decimal + assert isinstance(cost, Decimal) + except Exception as e: + # Expected if Anthropic API key is not set + if "ANTHROPIC_API_KEY" in str(e) or "API" in str(e): + pytest.skip("Anthropic API key not available for testing") + else: + raise + + def test_bedrock_cached_token_cost(self): + """Test cached token cost calculation for Bedrock Claude 3.5 Sonnet v2""" + model = "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0" + + # Test cached input token cost: 3e-07 per token + cached_cost = calculate_cost_by_tokens(1000, model, "cached") + assert cached_cost == Decimal("0.0003") + + def test_invalid_bedrock_model(self): + """Test that invalid Bedrock model names raise appropriate errors""" + model = "bedrock/invalid-model-name" + + with pytest.raises(KeyError) as exc_info: + calculate_cost_by_tokens(1000, model, "input") + + assert "not implemented" in str(exc_info.value).lower() \ No newline at end of file diff --git a/tokencost/costs.py b/tokencost/costs.py index ebb9756..14204a9 100644 --- a/tokencost/costs.py +++ b/tokencost/costs.py @@ -82,6 +82,17 @@ def strip_ft_model_name(model: str) -> str: return model +def normalize_bedrock_model_name(model: str) -> str: + """ + Normalize AWS Bedrock model names by removing the 'bedrock/' prefix. + Bedrock models format: bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0 + The cost dictionary uses the format without 'bedrock/' prefix. + """ + if model.startswith("bedrock/"): + model = model[8:] # Remove 'bedrock/' prefix + return model + + def count_message_tokens(messages: List[Dict[str, str]], model: str) -> int: """ Return the total number of tokens in a prompt's messages. @@ -95,6 +106,7 @@ def count_message_tokens(messages: List[Dict[str, str]], model: str) -> int: """ model = model.lower() model = strip_ft_model_name(model) + model = normalize_bedrock_model_name(model) # Anthropic token counting requires a valid API key if "claude-" in model: @@ -169,6 +181,7 @@ def count_string_tokens(prompt: str, model: str) -> int: int: The number of tokens in the text string. """ model = model.lower() + model = normalize_bedrock_model_name(model) if "/" in model: model = model.split("/")[-1] @@ -200,6 +213,7 @@ def calculate_cost_by_tokens(num_tokens: int, model: str, token_type: TokenType) Decimal: The calculated cost in USD. """ model = model.lower() + model = normalize_bedrock_model_name(model) if model not in TOKEN_COSTS: raise KeyError( f"""Model {model} is not implemented. @@ -238,6 +252,7 @@ def calculate_prompt_cost(prompt: Union[List[dict], str], model: str) -> Decimal """ model = model.lower() model = strip_ft_model_name(model) + model = normalize_bedrock_model_name(model) if model not in TOKEN_COSTS: raise KeyError( f"""Model {model} is not implemented. @@ -272,7 +287,9 @@ def calculate_completion_cost(completion: str, model: str) -> Decimal: >>> calculate_completion_cost(completion, "gpt-3.5-turbo") Decimal('0.000014') """ + model = model.lower() model = strip_ft_model_name(model) + model = normalize_bedrock_model_name(model) if model not in TOKEN_COSTS: raise KeyError( f"""Model {model} is not implemented. diff --git a/tokencost/model_prices.json b/tokencost/model_prices.json index 0f6f4b8..ff66562 100644 --- a/tokencost/model_prices.json +++ b/tokencost/model_prices.json @@ -4873,6 +4873,20 @@ "supports_pdf_input": true, "supports_tool_choice": true }, + "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0": { + "max_tokens": 4096, + "max_input_tokens": 200000, + "max_output_tokens": 4096, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_tool_choice": true + }, "anthropic.claude-3-5-sonnet-20241022-v2:0": { "supports_computer_use": true, "max_tokens": 8192, @@ -4892,6 +4906,25 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0": { + "supports_computer_use": true, + "max_tokens": 8192, + "max_input_tokens": 200000, + "max_output_tokens": 8192, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "bedrock", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "supports_pdf_input": true, + "supports_assistant_prefill": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "anthropic.claude-3-5-sonnet-latest-v2:0": { "max_tokens": 4096, "max_input_tokens": 200000,