fix(vertex_ai): bill cache creation tokens at the cache rate in above-128k pricing

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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)