diff --git a/.gitignore b/.gitignore index eb41980..c98e1e6 100644 --- a/.gitignore +++ b/.gitignore @@ -164,3 +164,4 @@ cython_debug/ .secrets .DS_STORE +/Flashpost-tests 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/model_prices.json b/tokencost/model_prices.json index 1ea705c..2d057a2 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,