fix(vertex_ai): keep Gemini cache creation token count out of request kwargs (#46152)

litellm/litellm_core_utils/litellm_logging.py

@@ -737,6 +737,7 @@ class Logging(LiteLLMLoggingBaseClass):
 
         # Init Caching related details
         self.caching_details: CachingDetails | None = None
+        self.vertex_ai_cache_creation_input_tokens: int = 0
         # Timing for results that cannot carry ``_hidden_params`` (plain-dict /v1/messages
         # responses and the bridge stream wrappers); see ``update_response_metadata``.
         self.response_timing_metrics: Mapping[str, float] = {}

litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py

@@ -16,7 +16,6 @@ from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
 )
 from litellm.llms.openai.openai import AllMessageValues
-from litellm.types.llms.vertex_ai import VERTEX_AI_CACHE_CREATION_TOKENS_KEY
 from litellm.utils import is_prompt_caching_valid_prompt
 
 from ..common_utils import VertexAIError, get_vertex_base_url
@@ -65,7 +64,7 @@ class _CreatedCachedContentRoot(RootModel[_CreatedCachedContent]):
 
 def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedContent) -> None:
     tokens: Final = (created.get("usageMetadata") or {}).get("totalTokenCount") or 0
-    logging_obj.model_call_details[VERTEX_AI_CACHE_CREATION_TOKENS_KEY] = tokens  # rebind-ok: shared kwargs bag
+    logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0)  # rebind-ok: set provider count
 
 
 class ContextCachingEndpoints(VertexBase):

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -54,7 +54,6 @@ from litellm.types.llms.openai import (
     OpenAIChatCompletionFinishReason,
 )
 from litellm.types.llms.vertex_ai import (
-    VERTEX_AI_CACHE_CREATION_TOKENS_KEY,
     VERTEX_AI_PROVIDER_METADATA_FIELDS,
     VERTEX_CREDENTIALS_TYPES,
     Candidates,
@@ -150,8 +149,8 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr
 
 
 def _cache_creation_tokens(logging_obj: LoggingClass) -> int:
-    recorded: Final = logging_obj.model_call_details.get(VERTEX_AI_CACHE_CREATION_TOKENS_KEY)
-    return recorded if isinstance(recorded, int) else 0
+    recorded: Final[object] = getattr(logging_obj, "vertex_ai_cache_creation_input_tokens", 0)
+    return recorded if isinstance(recorded, int) and recorded > 0 else 0
 
 
 def _served_model_name(model_version: object) -> str | None:

litellm/types/llms/vertex_ai.py

@@ -583,9 +583,6 @@ class VertexAICachedContentResponseObject(TypedDict):
     model: str
 
 
-VERTEX_AI_CACHE_CREATION_TOKENS_KEY: Final = "vertex_ai_cache_creation_input_tokens"
-
-
 class TaskTypeEnum(Enum):
     TASK_TYPE_UNSPECIFIED = "TASK_TYPE_UNSPECIFIED"
     RETRIEVAL_QUERY = "RETRIEVAL_QUERY"

tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py

@@ -2316,8 +2316,6 @@ def _cacheable_messages() -> list:
 @patch.object(ContextCachingEndpoints, "check_cache", return_value=None)
 @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u"))
 def test_check_and_create_cache_records_creation_tokens_for_billing(mock_url, mock_check):
-    from litellm.types.llms.vertex_ai import VERTEX_AI_CACHE_CREATION_TOKENS_KEY
-
     logging_obj = MagicMock()
     logging_obj.model_call_details = {}
     client = MagicMock(spec=HTTPHandler)
@@ -2339,16 +2337,15 @@ def test_check_and_create_cache_records_creation_tokens_for_billing(mock_url, mo
     )
 
     assert cache_name == "cachedContents/created"
-    assert logging_obj.model_call_details[VERTEX_AI_CACHE_CREATION_TOKENS_KEY] == 1280
+    assert logging_obj.vertex_ai_cache_creation_input_tokens == 1280
 
 
 @patch.object(ContextCachingEndpoints, "check_cache", return_value="cachedContents/existing")
 @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u"))
 def test_check_and_create_cache_records_no_creation_tokens_when_cache_exists(mock_url, mock_check):
-    from litellm.types.llms.vertex_ai import VERTEX_AI_CACHE_CREATION_TOKENS_KEY
-
     logging_obj = MagicMock()
     logging_obj.model_call_details = {}
+    logging_obj.vertex_ai_cache_creation_input_tokens = 0
     client = MagicMock(spec=HTTPHandler)
 
     _, _, cache_name = ContextCachingEndpoints().check_and_create_cache(
@@ -2368,15 +2365,13 @@ def test_check_and_create_cache_records_no_creation_tokens_when_cache_exists(moc
 
     assert cache_name == "cachedContents/existing"
     client.post.assert_not_called()
-    assert VERTEX_AI_CACHE_CREATION_TOKENS_KEY not in logging_obj.model_call_details
+    assert logging_obj.vertex_ai_cache_creation_input_tokens == 0
 
 
 @pytest.mark.asyncio
 @patch.object(ContextCachingEndpoints, "async_check_cache", return_value=None)
 @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u"))
 async def test_async_check_and_create_cache_records_creation_tokens_for_billing(mock_url, mock_check):
-    from litellm.types.llms.vertex_ai import VERTEX_AI_CACHE_CREATION_TOKENS_KEY
-
     logging_obj = MagicMock()
     logging_obj.model_call_details = {}
     client = MagicMock(spec=AsyncHTTPHandler)
@@ -2398,4 +2393,33 @@ async def test_async_check_and_create_cache_records_creation_tokens_for_billing(
     )
 
     assert cache_name == "cachedContents/created"
-    assert logging_obj.model_call_details[VERTEX_AI_CACHE_CREATION_TOKENS_KEY] == 2048
+    assert logging_obj.vertex_ai_cache_creation_input_tokens == 2048
+
+
+def test_check_and_create_cache_clamps_negative_creation_tokens() -> None:
+    logging_obj = MagicMock(spec=Logging)
+    logging_obj.model_call_details = {}
+    client = MagicMock(spec=HTTPHandler)
+    client.post.return_value = MagicMock(json=lambda: _created_cache_body(-500))
+
+    with (
+        patch.object(ContextCachingEndpoints, "check_cache", return_value=None),
+        patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching", return_value=("token", "https://u")),
+    ):
+        _, _, cache_name = ContextCachingEndpoints().check_and_create_cache(
+            messages=_cacheable_messages(),
+            optional_params={},
+            api_key="k",
+            api_base=None,
+            model="gemini-3.1-pro",
+            client=client,
+            timeout=30.0,
+            logging_obj=logging_obj,
+            custom_llm_provider="gemini",
+            vertex_project=None,
+            vertex_location=None,
+            vertex_auth_header=None,
+        )
+
+    assert cache_name == "cachedContents/created"
+    assert logging_obj.vertex_ai_cache_creation_input_tokens == 0

tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py

@@ -9515,6 +9515,39 @@ def test_vertex_ai_gemini_audio_ogg():
     print(response)
 
 
+def test_vertex_ai_cache_creation_kwarg_cannot_change_usage_or_cost() -> None:
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+    from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.headers = {"Content-Type": "application/json"}
+    mock_response.json.return_value = {
+        "candidates": [
+            {"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP", "index": 0}
+        ],
+        "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15},
+    }
+    client = HTTPHandler()
+    httpx_mock = MagicMock(return_value=mock_response)
+    with (
+        patch.object(client, "post", new=httpx_mock),
+        patch.object(VertexBase, "_ensure_access_token", return_value=("fake-token", "fake-project")),
+    ):
+        response = litellm.completion(
+            model="vertex_ai/gemini-2.5-flash",
+            messages=[{"role": "user", "content": "hi"}],
+            vertex_project="fake-project",
+            vertex_location="us-central1",
+            vertex_ai_cache_creation_input_tokens=-1000000,
+            client=client,
+        )
+
+    assert response.usage.prompt_tokens == 10
+    assert response.usage.cache_creation_input_tokens is None
+    assert response._hidden_params["response_cost"] > 0
+
+
 def load_vertex_ai_credentials():
     # Define the path to the vertex_key.json file
     print("loading vertex ai credentials")
@@ -9637,10 +9670,9 @@ def test_calculate_usage_without_cache_creation_leaves_usage_unchanged():
 
 def test_streaming_usage_reads_cache_creation_tokens_recorded_on_logging_obj():
     from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ModelResponseIterator
-    from litellm.types.llms.vertex_ai import VERTEX_AI_CACHE_CREATION_TOKENS_KEY
 
     logging_obj = MagicMock()
-    logging_obj.model_call_details = {VERTEX_AI_CACHE_CREATION_TOKENS_KEY: 1280}
+    logging_obj.vertex_ai_cache_creation_input_tokens = 1280
     iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=logging_obj)
 
     chunk = iterator.chunk_parser(