From d39b2e8888820966574ffc80a301be4e9fffe485 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 19 Aug 2025 16:56:13 +0530 Subject: [PATCH] Add test for long context cost calculation --- tests/test_litellm/test_cost_calculator.py | 119 +++++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 4efba133f1..e229effb5f 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -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}")