From 49ed3065f6905dfaa2676a313ee63ac2cb766187 Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 30 Jun 2025 22:00:10 -0700
Subject: [PATCH] VertexAI Anthropic - streaming cost tracking w/ prompt
 caching fixes (#12188)

* fix(rebuild-usage-object---ensure-cache_tokens-is-set): Ensures cache tokens is correctly set

Fixes https://github.com/BerriAI/litellm/issues/12149

* test(test_stream_chunk_builder_utils.py): add unit test to ensure cached tokens is part of stream chunk builder

Ensures standardized values are used
---
 .../streaming_chunk_builder_utils.py          |  8 +++--
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 32 +++++++++++++++++++
 .../test_streaming_chunk_builder_utils.py     |  1 +
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py
index 4068d2e043c..0517d27e299 100644
--- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py
+++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py
@@ -107,9 +107,9 @@ class ChunkProcessor:
         self, tool_call_chunks: List[Dict[str, Any]]
     ) -> List[ChatCompletionMessageToolCall]:
         tool_calls_list: List[ChatCompletionMessageToolCall] = []
-        tool_call_map: Dict[
-            int, Dict[str, Any]
-        ] = {}  # Map to store tool calls by index
+        tool_call_map: Dict[int, Dict[str, Any]] = (
+            {}
+        )  # Map to store tool calls by index
 
         for chunk in tool_call_chunks:
             choices = chunk["choices"]
@@ -415,6 +415,8 @@ class ChunkProcessor:
         if prompt_tokens_details is not None:
             returned_usage.prompt_tokens_details = prompt_tokens_details
 
+        # Return a new usage object with the new values
+        returned_usage = Usage(**returned_usage.model_dump())
         return returned_usage
 
 
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index 7df783e719d..c988a99f851 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -140,3 +140,35 @@ def test_generic_cost_per_token_above_200k_tokens():
         * usage.completion_tokens,
         10,
     )
+
+
+def test_generic_cost_per_token_anthropic_prompt_caching():
+    model = "claude-sonnet-4@20250514"
+    usage = Usage(
+        completion_tokens=90,
+        prompt_tokens=28436,
+        total_tokens=28526,
+        completion_tokens_details=CompletionTokensDetailsWrapper(
+            accepted_prediction_tokens=None,
+            audio_tokens=None,
+            reasoning_tokens=0,
+            rejected_prediction_tokens=None,
+            text_tokens=None,
+        ),
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None
+        ),
+        cache_creation_input_tokens=118,
+        cache_read_input_tokens=28432,
+    )
+
+    custom_llm_provider = "vertex_ai"
+
+    prompt_cost, completion_cost = generic_cost_per_token(
+        model=model,
+        usage=usage,
+        custom_llm_provider=custom_llm_provider,
+    )
+
+    print(f"prompt_cost: {prompt_cost}")
+    assert prompt_cost < 0.085
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
index dfc3f01d112..facb5cef3cd 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
@@ -242,3 +242,4 @@ def test_cache_read_input_tokens_retained():
 
     assert usage.cache_creation_input_tokens == 4
     assert usage.cache_read_input_tokens == 11775
+    assert usage.prompt_tokens_details.cached_tokens == 11775