Add test for long context cost calculation

2026-07-10 19:04:29 +00:00 · 2025-08-19 16:56:13 +05:30
parent 4f38a63152
commit d39b2e8888
1 changed files with 119 additions and 0 deletions
@@ -480,3 +480,122 @@ def test_gemini_25_implicit_caching_cost():
    ), f"Expected cost {expected_cost}, but got {result}"

    print(f"✓ Gemini 2.5 implicit caching cost calculation is correct: ${result:.8f}")
+
+
+def test_log_context_cost_calculation():
+    """
+    Test that log context cost calculation works correctly with tiered pricing.
+    
+    This test verifies that when using extended context (above 200k tokens),
+    the log context costs are calculated using the appropriate tiered rates.
+    """
+    from litellm import completion_cost
+    from litellm.types.utils import (
+        Choices,
+        Message,
+        ModelResponse,
+        PromptTokensDetailsWrapper,
+        Usage,
+    )
+
+    # Create a mock response with extended context usage
+    extended_context_response = ModelResponse(
+        id="test-extended-context-response",
+        created=1750733889,
+        model="claude-4-sonnet-20250514",
+        object="chat.completion",
+        system_fingerprint=None,
+        choices=[
+            Choices(
+                finish_reason="stop",
+                index=0,
+                message=Message(
+                    content="This is a test response for extended context cost calculation.",
+                    role="assistant",
+                    tool_calls=None,
+                    function_call=None,
+                ),
+            )
+        ],
+        usage=Usage(
+            total_tokens=350000,  # Above 200k threshold
+            prompt_tokens=300000,  # Above 200k threshold
+            completion_tokens=50000,
+            prompt_tokens_details=PromptTokensDetailsWrapper(
+                text_tokens=300000,
+                cached_tokens=0,  # No cache hits
+                audio_tokens=None,
+                image_tokens=None,
+                character_count=None,
+                video_length_seconds=None,
+            ),
+            completion_tokens_details=None,
+            _cache_creation_input_tokens=1000,  # Some tokens added to cache
+        ),
+    )
+
+    # Calculate the cost using the extended context model
+    result = completion_cost(
+        completion_response=extended_context_response,
+        model="claude-4-sonnet-20250514",
+        custom_llm_provider="anthropic",
+    )
+
+    # Debug: Print the actual result
+    print(f"DEBUG: Actual cost result: ${result:.6f}")
+    
+    # Get model info to understand the pricing
+    from litellm import get_model_info
+    model_info = get_model_info(model="claude-4-sonnet-20250514", custom_llm_provider="anthropic")
+    
+    # Calculate expected cost based on actual model pricing
+    input_cost_per_token = model_info.get("input_cost_per_token", 0)
+    output_cost_per_token = model_info.get("output_cost_per_token", 0)
+    cache_creation_cost_per_token = model_info.get("cache_creation_input_token_cost", 0)
+    
+    # Check if tiered pricing is applied
+    input_cost_above_200k = model_info.get("input_cost_per_token_above_200k_tokens", input_cost_per_token)
+    output_cost_above_200k = model_info.get("output_cost_per_token_above_200k_tokens", output_cost_per_token)
+    cache_creation_above_200k = model_info.get("cache_creation_input_token_cost_above_200k_tokens", cache_creation_cost_per_token)
+    
+    print(f"DEBUG: Base input cost per token: ${input_cost_per_token:.2e}")
+    print(f"DEBUG: Base output cost per token: ${output_cost_per_token:.2e}")
+    print(f"DEBUG: Base cache creation cost per token: ${cache_creation_cost_per_token:.2e}")
+    
+    # Handle tiered pricing - if not available, use base pricing
+    if input_cost_above_200k is not None:
+        print(f"DEBUG: Tiered input cost per token (>200k): ${input_cost_above_200k:.2e}")
+    else:
+        print(f"DEBUG: No tiered input pricing available, using base pricing")
+        input_cost_above_200k = input_cost_per_token
+        
+    if output_cost_above_200k is not None:
+        print(f"DEBUG: Tiered output cost per token (>200k): ${output_cost_above_200k:.2e}")
+    else:
+        print(f"DEBUG: No tiered output pricing available, using base pricing")
+        output_cost_above_200k = output_cost_per_token
+        
+    if cache_creation_above_200k is not None:
+        print(f"DEBUG: Tiered cache creation cost per token (>200k): ${cache_creation_above_200k:.2e}")
+    else:
+        print(f"DEBUG: No tiered cache creation pricing available, using base pricing")
+        cache_creation_above_200k = cache_creation_cost_per_token
+    
+    # Since we're above 200k tokens, we should use tiered pricing if available
+    expected_input_cost = 300000 * input_cost_above_200k
+    expected_output_cost = 50000 * output_cost_above_200k
+    expected_cache_cost = 1000 * cache_creation_above_200k
+    expected_total = expected_input_cost + expected_output_cost + expected_cache_cost
+    
+    print(f"DEBUG: Expected total: ${expected_total:.6f}")
+    
+    # Allow for small floating point differences
+    assert (
+        abs(result - expected_total) < 1e-6
+    ), f"Expected cost ${expected_total:.6f}, but got ${result:.6f}"
+
+    print(f"✓ Log context cost calculation with tiered pricing is correct: ${result:.6f}")
+    print(f"  - Input tokens (300k): ${expected_input_cost:.6f}")
+    print(f"  - Output tokens (50k): ${expected_output_cost:.6f}")
+    print(f"  - Cache creation (1k): ${expected_cache_cost:.6f}")
+    print(f"  - Total: ${result:.6f}")