diff --git a/litellm/llms/vertex_ai/cost_calculator.py b/litellm/llms/vertex_ai/cost_calculator.py index 8b00fc2e925..51c9434660c 100644 --- a/litellm/llms/vertex_ai/cost_calculator.py +++ b/litellm/llms/vertex_ai/cost_calculator.py @@ -183,13 +183,30 @@ def _handle_128k_pricing( input_cost_per_token_above_128k_tokens: Final = model_info.get("input_cost_per_token_above_128k_tokens") output_cost_per_token_above_128k_tokens = model_info.get("output_cost_per_token_above_128k_tokens") - prompt_tokens: Final = usage.prompt_tokens + prompt_tokens_details: Final = usage.prompt_tokens_details + cache_read_tokens: Final = ( + (prompt_tokens_details.cached_tokens or 0) if prompt_tokens_details is not None else 0 + ) + cache_creation_tokens: Final = ( + (prompt_tokens_details.cache_creation_tokens or 0) if prompt_tokens_details is not None else 0 + ) + text_tokens: Final = max(usage.prompt_tokens - cache_read_tokens - cache_creation_tokens, 0) completion_tokens: Final = usage.completion_tokens - if _is_above_128k(tokens=prompt_tokens) and input_cost_per_token_above_128k_tokens is not None: - prompt_cost = prompt_tokens * input_cost_per_token_above_128k_tokens - else: - prompt_cost = prompt_tokens * (model_info["input_cost_per_token"] or 0.0) + input_rate: Final = ( + input_cost_per_token_above_128k_tokens + if input_cost_per_token_above_128k_tokens is not None and _is_above_128k(tokens=text_tokens) + else (model_info["input_cost_per_token"] or 0.0) + ) + + cache_read_rate: Final = model_info.get("cache_read_input_token_cost") or input_rate + cache_creation_rate: Final = model_info.get("cache_creation_input_token_cost") or input_rate + + prompt_cost = ( + text_tokens * input_rate + + cache_read_tokens * cache_read_rate + + cache_creation_tokens * cache_creation_rate + ) ## CALCULATE OUTPUT COST output_cost_per_token_above_128k_tokens = model_info.get("output_cost_per_token_above_128k_tokens") diff --git a/tests/test_litellm/llms/vertex_ai/test_cost_calculator.py b/tests/test_litellm/llms/vertex_ai/test_cost_calculator.py new file mode 100644 index 00000000000..8e6e49f4ca2 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_cost_calculator.py @@ -0,0 +1,48 @@ +from typing import Final + +import pytest + +import litellm +from litellm.llms.vertex_ai.cost_calculator import cost_per_token +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + +@pytest.mark.parametrize( + ("text_tokens", "expected_prompt_cost"), + [ + (140_000, 140_000 * 0.002 + 120_000 * 0.0005), + (120_000, 120_000 * 0.001 + 120_000 * 0.0005), + ], + ids=["creation_tokens_do_not_change_the_tier_rate", "creation_tokens_cannot_push_the_tier_threshold"], +) +def test_above_128k_pricing_splits_cache_creation_tokens_out_of_the_prompt( + monkeypatch: pytest.MonkeyPatch, text_tokens: int, expected_prompt_cost: float +) -> None: + """Cache creation tokens bill at the cache-creation rate and never count toward the above-128k tier.""" + model: Final = "vertex_ai/fake-above-128k-model" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "vertex_ai", + "input_cost_per_token": 0.001, + "input_cost_per_token_above_128k_tokens": 0.002, + "cache_creation_input_token_cost": 0.0005, + "output_cost_per_token": 0.003, + }, + ) + + usage: Final = Usage( + prompt_tokens=text_tokens + 120_000, + completion_tokens=10, + total_tokens=text_tokens + 120_010, + prompt_tokens_details=PromptTokensDetailsWrapper(cache_creation_tokens=120_000), + ) + prompt_cost, completion_cost = cost_per_token( + model=model, + custom_llm_provider="vertex_ai", + usage=usage, + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + assert completion_cost == pytest.approx(10 * 0.003)