From fdab308646648c246d2826e621d35b41b4ce3b49 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:19:12 -0700 Subject: [PATCH] fix(vertex_ai): keep Gemini cache creation token count out of request kwargs (#46152) --- litellm/litellm_core_utils/litellm_logging.py | 1 + .../vertex_ai_context_caching.py | 3 +- .../vertex_and_google_ai_studio_gemini.py | 5 +-- litellm/types/llms/vertex_ai.py | 3 -- .../test_vertex_ai_context_caching.py | 42 +++++++++++++++---- ...test_vertex_and_google_ai_studio_gemini.py | 36 +++++++++++++++- 6 files changed, 71 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5395383f8bd..d73d8813308 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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] = {} diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index c8608705eff..06b8f3a3fbb 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/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): diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index e3b43e2bb64..517ac9e98b1 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/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: diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 28a174729a8..116fdda0b5d 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/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" diff --git a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index abc0ec661a0..19e7e2dee52 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/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 diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index bc885891206..3f49869b19f 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/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(