From ddd471ffa6de7242a22abfd4dbcbc7c5af536053 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 23 Sep 2026 02:58:33 +0000 Subject: [PATCH] fix(vertex_ai): bill cache creation at the resolved input tier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/vertex_ai/cost_calculator.py | 2 +- .../vertex_ai/context_caching/__init__.py | 1 + .../test_vertex_ai_context_caching.py | 119 ++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 30 +++++ .../llms/vertex_ai/test_cost_calculator.py | 31 +++++ 5 files changed, 182 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/llms/vertex_ai/context_caching/__init__.py create mode 100644 tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py diff --git a/litellm/llms/vertex_ai/cost_calculator.py b/litellm/llms/vertex_ai/cost_calculator.py index e0a79e00ab3..18ef7b99e98 100644 --- a/litellm/llms/vertex_ai/cost_calculator.py +++ b/litellm/llms/vertex_ai/cost_calculator.py @@ -197,7 +197,7 @@ def _handle_128k_pricing( input_rate: Final = input_cost_per_token_above_128k_tokens if above_tier else base_input_rate 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 base_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 diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/__init__.py b/tests/test_litellm/llms/vertex_ai/context_caching/__init__.py new file mode 100644 index 00000000000..09fd91a341c --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/context_caching/__init__.py @@ -0,0 +1 @@ +"""Vertex AI context caching tests package.""" diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py new file mode 100644 index 00000000000..5cc89189960 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -0,0 +1,119 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + ContextCachingEndpoints, +) + +_CACHE_RESPONSE: Final = { + "name": "cachedContents/new-cache-name", + "model": "gemini-2.5-flash", + "usageMetadata": {"totalTokenCount": 2048}, + "createTime": "2026-09-01T00:00:00Z", + "expireTime": "2026-09-01T01:00:00Z", +} + +_EXPECTED_CREATION: Final = { + "name": "cachedContents/new-cache-name", + "model": "gemini-2.5-flash", + "total_token_count": 2048, + "create_time": "2026-09-01T00:00:00Z", + "expire_time": "2026-09-01T01:00:00Z", +} + + +@patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" +) +@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") +@patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" +) +@patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt" +) +@patch.object(ContextCachingEndpoints, "check_cache") +@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") +def test_check_and_create_cache_stores_cached_content_creation_on_logging( + mock_get_token_url, mock_check_cache, mock_valid_prompt, mock_transform, mock_cache_obj, mock_separate +) -> None: + mock_separate.return_value = ([{"role": "user", "content": "cached"}], [{"role": "user", "content": "fresh"}]) + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = None + mock_valid_prompt.return_value = True + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-2.5-flash", "contents": []} + response: Final = MagicMock() + response.json.return_value = _CACHE_RESPONSE + client: Final = MagicMock(spec=HTTPHandler) + client.post.return_value = response + logging_obj: Final = MagicMock(spec=Logging) + logging_obj.model_call_details = {} + + ContextCachingEndpoints().check_and_create_cache( + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + api_key="test_key", + api_base=None, + model="gemini-2.5-flash", + client=client, + timeout=30.0, + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="vertex_test_token", + ) + + assert logging_obj.model_call_details["vertex_ai_cached_content"] == _EXPECTED_CREATION + + +@pytest.mark.asyncio +@patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" +) +@patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") +@patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" +) +@patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt" +) +@patch.object(ContextCachingEndpoints, "async_check_cache") +@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") +async def test_async_check_and_create_cache_stores_cached_content_creation_on_logging( + mock_get_token_url, mock_async_check_cache, mock_valid_prompt, mock_transform, mock_cache_obj, mock_separate +) -> None: + mock_separate.return_value = ([{"role": "user", "content": "cached"}], [{"role": "user", "content": "fresh"}]) + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_async_check_cache.return_value = None + mock_valid_prompt.return_value = True + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-2.5-flash", "contents": []} + response: Final = MagicMock() + response.json.return_value = _CACHE_RESPONSE + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.post = AsyncMock(return_value=response) + logging_obj: Final = MagicMock(spec=Logging) + logging_obj.model_call_details = {} + + await ContextCachingEndpoints().async_check_and_create_cache( + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + api_key="test_key", + api_base=None, + model="gemini-2.5-flash", + client=client, + timeout=30.0, + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="vertex_test_token", + ) + + assert logging_obj.model_call_details["vertex_ai_cached_content"] == _EXPECTED_CREATION diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 35c827fff2a..ffaf7366ec0 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -6440,3 +6440,33 @@ def test_transform_response_applies_cache_creation_usage(include_creation): assert result.usage.prompt_tokens == 10010 assert not hasattr(result.usage, "cache_creation_input_tokens") assert "vertex_ai_cached_content" not in result._hidden_params + + +def test_streaming_chunk_parser_surfaces_cache_creation_usage_and_metadata(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ModelResponseIterator + + model = "gemini-3.8-flash" + cached_content = { + "name": "cached-content", + "model": model, + "total_token_count": 10000, + "create_time": None, + "expire_time": None, + } + litellm_logging = MagicMock() + litellm_logging.model_call_details = {"vertex_ai_cached_content": cached_content} + chunk = { + "candidates": [{"content": {"parts": [{"text": "Hello"}]}}], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 1, + "totalTokenCount": 11, + }, + } + + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) + streaming_chunk = iterator.chunk_parser(chunk) + + assert streaming_chunk.usage.prompt_tokens == 10010 + assert streaming_chunk.usage.cache_creation_input_tokens == 10000 + assert streaming_chunk._hidden_params["vertex_ai_cached_content"] == cached_content diff --git a/tests/test_litellm/llms/vertex_ai/test_cost_calculator.py b/tests/test_litellm/llms/vertex_ai/test_cost_calculator.py index 7c558c3f549..4f263e90aeb 100644 --- a/tests/test_litellm/llms/vertex_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/vertex_ai/test_cost_calculator.py @@ -61,3 +61,34 @@ def test_above_128k_pricing_splits_cache_tokens_out_of_the_prompt( assert prompt_cost == pytest.approx(expected_prompt_cost) assert completion_cost == pytest.approx(10 * 0.003) + + +def test_above_128k_pricing_falls_back_to_the_resolved_tier_rate_for_creation_tokens( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A missing cache-creation rate resolves to the tiered input rate, like generic_cost_per_token.""" + model: Final = "vertex_ai/fake-above-128k-model-no-creation-rate" + 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, + "output_cost_per_token": 0.003, + }, + ) + + usage: Final = Usage( + prompt_tokens=260_000, + completion_tokens=10, + total_tokens=260_010, + prompt_tokens_details=PromptTokensDetailsWrapper(cache_creation_tokens=120_000), + ) + prompt_cost, _ = cost_per_token( + model=model, + custom_llm_provider="vertex_ai", + usage=usage, + ) + + assert prompt_cost == pytest.approx(140_000 * 0.002 + 120_000 * 0.002)