Add aws bedrock model cost tracking

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.

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()

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.

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,