From 92e8e98750adddf68d59413ba9eb1264a3e6c367 Mon Sep 17 00:00:00 2001
From: Genmin <joey@joeyroth.com>
Date: Thu, 30 Apr 2026 09:56:36 -0700
Subject: [PATCH] fix: apply cached token rate for custom pricing

---
 litellm/cost_calculator.py                 | 22 +++++++-
 tests/test_litellm/test_cost_calculator.py | 65 ++++++++++++++++++++++
 2 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 8a68d74be5b..65c09940b04 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -172,6 +172,8 @@ _MCP_CALL_TYPE = CallTypes.call_mcp_tool.value
 def _cost_per_token_custom_pricing_helper(
     prompt_tokens: float = 0,
     completion_tokens: float = 0,
+    cache_read_input_tokens: Optional[int] = 0,
+    usage_object: Optional[Usage] = None,
     response_time_ms: Optional[float] = 0.0,
     ### CUSTOM PRICING ###
     custom_cost_per_token: Optional[CostPerToken] = None,
@@ -182,7 +184,23 @@ def _cost_per_token_custom_pricing_helper(
         return None
 
     if custom_cost_per_token is not None:
-        input_cost = custom_cost_per_token["input_cost_per_token"] * prompt_tokens
+        cache_read_input_token_cost = cast(
+            Optional[float],
+            custom_cost_per_token.get("cache_read_input_token_cost"),
+        )
+        if cache_read_input_token_cost is not None:
+            if not cache_read_input_tokens and usage_object is not None:
+                cache_read_input_tokens = _parse_prompt_tokens_details(usage_object)[
+                    "cache_hit_tokens"
+                ]
+            cache_read_tokens = max(0, cache_read_input_tokens or 0)
+            uncached_prompt_tokens = max(0, prompt_tokens - cache_read_tokens)
+            input_cost = (
+                custom_cost_per_token["input_cost_per_token"] * uncached_prompt_tokens
+                + cache_read_input_token_cost * cache_read_tokens
+            )
+        else:
+            input_cost = custom_cost_per_token["input_cost_per_token"] * prompt_tokens
         output_cost = custom_cost_per_token["output_cost_per_token"] * completion_tokens
         return input_cost, output_cost
     elif custom_cost_per_second is not None:
@@ -326,6 +344,8 @@ def cost_per_token(  # noqa: PLR0915
     response_cost = _cost_per_token_custom_pricing_helper(
         prompt_tokens=prompt_tokens,
         completion_tokens=completion_tokens,
+        cache_read_input_tokens=cache_read_input_tokens,
+        usage_object=usage_block,
         response_time_ms=response_time_ms,
         custom_cost_per_second=custom_cost_per_second,
         custom_cost_per_token=custom_cost_per_token,
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index ebe175b2503..02881e2af1b 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -48,6 +48,71 @@ def test_completion_cost_uses_response_model_for_dynamic_routing():
     assert cost > 0, "Cost should be calculated using response model"
 
 
+def test_completion_cost_custom_pricing_uses_cache_read_rate():
+    usage = Usage(
+        prompt_tokens=6074,
+        completion_tokens=285,
+        total_tokens=6359,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=3456,
+            audio_tokens=0,
+        ),
+    )
+    response = ModelResponse(
+        id="test-id",
+        created=1234567890,
+        model="openai/gpt-5.4",
+        object="chat.completion",
+        choices=[],
+        usage=usage,
+    )
+
+    cost = completion_cost(
+        completion_response=response,
+        model="openai/gpt-5.4",
+        custom_llm_provider="openai",
+        custom_cost_per_token={
+            "input_cost_per_token": 0.0000025,
+            "output_cost_per_token": 0.000015,
+            "cache_read_input_token_cost": 0.00000025,
+        },
+    )
+
+    assert cost == pytest.approx(0.011684)
+
+
+def test_completion_cost_custom_pricing_without_cache_read_rate_preserves_input_rate():
+    usage = Usage(
+        prompt_tokens=6074,
+        completion_tokens=285,
+        total_tokens=6359,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=3456,
+            audio_tokens=0,
+        ),
+    )
+    response = ModelResponse(
+        id="test-id",
+        created=1234567890,
+        model="openai/gpt-5.4",
+        object="chat.completion",
+        choices=[],
+        usage=usage,
+    )
+
+    cost = completion_cost(
+        completion_response=response,
+        model="openai/gpt-5.4",
+        custom_llm_provider="openai",
+        custom_cost_per_token={
+            "input_cost_per_token": 0.0000025,
+            "output_cost_per_token": 0.000015,
+        },
+    )
+
+    assert cost == pytest.approx(0.01946)
+
+
 def test_cost_calculator_with_response_cost_in_additional_headers():
     class MockResponse(BaseModel):
         _hidden_params = {