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 75d4ffbed86..049336905ef 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 @@ -15,6 +15,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.llms.openai.openai import AllMessageValues from litellm.types.llms.vertex_ai import ( CachedContentListAllResponseBody, + VertexAICachedContentCreation, VertexAICachedContentResponseObject, ) from litellm.utils import is_prompt_caching_valid_prompt @@ -30,6 +31,7 @@ from .transformation import ( local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for calling 'get_cache_key' function MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination +VERTEX_AI_CACHED_CONTENT_KEY: Final = "vertex_ai_cached_content" class ContextCachingEndpoints(VertexBase): @@ -424,9 +426,16 @@ class ContextCachingEndpoints(VertexBase): raise VertexAIError(status_code=408, message="Timeout error occurred.") raw_response_cached: Final = response.json() - cached_content_response_obj: Final = VertexAICachedContentResponseObject( - name=raw_response_cached.get("name"), model=raw_response_cached.get("model") + cached_content_response_obj: Final = VertexAICachedContentResponseObject(**raw_response_cached) + usage_metadata: Final = cached_content_response_obj.get("usageMetadata", {}) + cached_content_creation: Final = VertexAICachedContentCreation( + name=cached_content_response_obj["name"], + model=cached_content_response_obj["model"], + total_token_count=usage_metadata.get("totalTokenCount", 0), + create_time=cached_content_response_obj.get("createTime"), + expire_time=cached_content_response_obj.get("expireTime"), ) + logging_obj.model_call_details[VERTEX_AI_CACHED_CONTENT_KEY] = cached_content_creation return ( non_cached_messages, optional_params, @@ -579,9 +588,16 @@ class ContextCachingEndpoints(VertexBase): raise VertexAIError(status_code=408, message="Timeout error occurred.") raw_response_cached: Final = response.json() - cached_content_response_obj: Final = VertexAICachedContentResponseObject( - name=raw_response_cached.get("name"), model=raw_response_cached.get("model") + cached_content_response_obj: Final = VertexAICachedContentResponseObject(**raw_response_cached) + usage_metadata: Final = cached_content_response_obj.get("usageMetadata", {}) + cached_content_creation: Final = VertexAICachedContentCreation( + name=cached_content_response_obj["name"], + model=cached_content_response_obj["model"], + total_token_count=usage_metadata.get("totalTokenCount", 0), + create_time=cached_content_response_obj.get("createTime"), + expire_time=cached_content_response_obj.get("expireTime"), ) + logging_obj.model_call_details[VERTEX_AI_CACHED_CONTENT_KEY] = cached_content_creation return ( non_cached_messages, optional_params, 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 204b1ddd5ea..1317ffcd3c1 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 @@ -66,6 +66,7 @@ from litellm.types.llms.vertex_ai import ( ToolConfig, Tools, UsageMetadata, + VertexAICachedContentCreation, VertexToolName, ) from litellm.types.utils import ( @@ -133,6 +134,30 @@ def _served_model_name(model_version: object) -> str | None: return model_version.split("@", 1)[0] +def _add_cache_creation_usage(usage: Usage, creation: VertexAICachedContentCreation) -> Usage: + creation_tokens: Final = creation["total_token_count"] + if creation_tokens <= 0: + return usage + + prompt_tokens_details: Final = ( + PromptTokensDetailsWrapper(**usage.prompt_tokens_details.model_dump()) + if usage.prompt_tokens_details is not None + else PromptTokensDetailsWrapper() + ) + cache_read_tokens: Final = getattr(usage, "_cache_read_input_tokens", 0) or 0 + return Usage( + prompt_tokens=usage.prompt_tokens + creation_tokens, + completion_tokens=usage.completion_tokens, + total_tokens=usage.total_tokens + creation_tokens, + prompt_tokens_details=prompt_tokens_details, + completion_tokens_details=usage.completion_tokens_details, + server_tool_use=getattr(usage, "server_tool_use", None), + cost=getattr(usage, "cost", None), + cache_creation_input_tokens=creation_tokens, + **({"cache_read_input_tokens": cache_read_tokens} if cache_read_tokens > 0 else {}), + ) + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -2459,7 +2484,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _, # cumulative_tool_call_index not needed in non-streaming ) = VertexGeminiConfig._process_candidates(_candidates, model_response, logging_obj.optional_params) - usage: Final = VertexGeminiConfig._calculate_usage(completion_response=completion_response) + base_usage: Final = VertexGeminiConfig._calculate_usage(completion_response=completion_response) + from ..context_caching.vertex_ai_context_caching import VERTEX_AI_CACHED_CONTENT_KEY + + cached_content_creation: Final = logging_obj.model_call_details.get(VERTEX_AI_CACHED_CONTENT_KEY) + usage: Final = ( + _add_cache_creation_usage(base_usage, cast(VertexAICachedContentCreation, cached_content_creation)) + if isinstance(cached_content_creation, dict) and "total_token_count" in cached_content_creation + else base_usage + ) VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata) @@ -3211,9 +3244,17 @@ class ModelResponseIterator: if "usageMetadata" not in processed_chunk: return None - usage: Final = VertexGeminiConfig._calculate_usage( + base_usage: Final = VertexGeminiConfig._calculate_usage( completion_response=processed_chunk, ) + from ..context_caching.vertex_ai_context_caching import VERTEX_AI_CACHED_CONTENT_KEY + + cached_content_creation: Final = self.logging_obj.model_call_details.get(VERTEX_AI_CACHED_CONTENT_KEY) + usage: Final = ( + _add_cache_creation_usage(base_usage, cast(VertexAICachedContentCreation, cached_content_creation)) + if isinstance(cached_content_creation, dict) and "total_token_count" in cached_content_creation + else base_usage + ) VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index ce51e46ef15..696e0e8b606 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -2,6 +2,8 @@ from enum import Enum from typing import Any, Final, Literal, Protocol from typing_extensions import ( + NotRequired, + ReadOnly, Required, TypedDict, ) @@ -567,9 +569,24 @@ class MultimodalPredictions(TypedDict): predictions: list[MultimodalPrediction] +class VertexAICachedContentUsageMetadata(TypedDict): + totalTokenCount: ReadOnly[int] + + class VertexAICachedContentResponseObject(TypedDict): - name: str - model: str + name: ReadOnly[str] + model: ReadOnly[str] + usageMetadata: NotRequired[ReadOnly[VertexAICachedContentUsageMetadata]] + createTime: NotRequired[ReadOnly[str]] + expireTime: NotRequired[ReadOnly[str]] + + +class VertexAICachedContentCreation(TypedDict): + name: ReadOnly[str] + model: ReadOnly[str] + total_token_count: ReadOnly[int] + create_time: ReadOnly[str | None] + expire_time: ReadOnly[str | None] class TaskTypeEnum(Enum): 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 8b0662de3f9..2d7d27213ca 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 @@ -15,11 +15,12 @@ from litellm.llms.anthropic.experimental_pass_through.messages import handler as from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.gemini import vertex_and_google_ai_studio_gemini as vertex_gemini_module from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) -from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata -from litellm.types.utils import ChoiceLogprobs, Usage +from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata, VertexAICachedContentCreation +from litellm.types.utils import ChoiceLogprobs, PromptTokensDetailsWrapper, Usage from litellm.utils import CustomStreamWrapper @@ -6340,3 +6341,94 @@ def test_gemini_multi_candidate_messages_do_not_share_state(): assert resp.choices[1].message.tool_calls is None assert getattr(resp.choices[1].message, "reasoning_content", None) is None assert resp.choices[1].provider_specific_fields["native_finish_reason"] == "STOP" + + +def test_add_cache_creation_usage_preserves_cache_read_and_bills_creation_tokens(): + base_usage = Usage( + prompt_tokens=10010, + completion_tokens=1, + total_tokens=10011, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=10000), + ) + creation: VertexAICachedContentCreation = { + "name": "cached-content", + "model": "gemini-2.0-flash", + "total_token_count": 10000, + "create_time": None, + "expire_time": None, + } + + merged_usage = vertex_gemini_module._add_cache_creation_usage(base_usage, creation) + + assert merged_usage.prompt_tokens == 20010 + assert merged_usage.total_tokens == 20011 + assert merged_usage.cache_creation_input_tokens == 10000 + assert merged_usage.prompt_tokens_details.cached_tokens == 10000 + assert merged_usage.prompt_tokens_details.cache_creation_tokens == 10000 + + model = "gemini-2.0-flash" + base_response = ModelResponse(model=model, usage=base_usage) + merged_response = ModelResponse(model=model, usage=merged_usage) + base_cost = litellm.completion_cost( + completion_response=base_response, + model=model, + custom_llm_provider="vertex_ai", + ) + merged_cost = litellm.completion_cost( + completion_response=merged_response, + model=model, + custom_llm_provider="vertex_ai", + ) + model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") + creation_rate = model_info.get("cache_creation_input_token_cost") or model_info["input_cost_per_token"] + + assert merged_cost > base_cost + assert merged_cost == pytest.approx(base_cost + 10000 * creation_rate) + + +@pytest.mark.parametrize("include_creation", [True, False]) +def test_transform_response_applies_cache_creation_usage(include_creation): + model = "gemini-2.0-flash" + logging_obj = MagicMock() + logging_obj.model_call_details = ( + { + "vertex_ai_cached_content": { + "name": "cached-content", + "model": model, + "total_token_count": 10000, + "create_time": None, + "expire_time": None, + } + } + if include_creation + else {} + ) + raw_response = MagicMock() + raw_response.json.return_value = { + "candidates": [{"content": {"parts": [{"text": "Hello"}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 10010, + "cachedContentTokenCount": 10000, + "candidatesTokenCount": 1, + "totalTokenCount": 10011, + }, + } + + result = VertexGeminiConfig().transform_response( + model=model, + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + if include_creation: + assert result.usage.prompt_tokens == 20010 + assert result.usage.cache_creation_input_tokens == 10000 + else: + assert result.usage.prompt_tokens == 10010 + assert not hasattr(result.usage, "cache_creation_input_tokens") 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 7913700c8a7..a06a9f49d02 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 @@ -1,4 +1,3 @@ -from typing import List from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -40,6 +39,7 @@ class TestContextCachingEndpoints: """Setup for each test method""" self.context_caching = ContextCachingEndpoints() self.mock_logging = MagicMock(spec=Logging) + self.mock_logging.model_call_details = {} self.mock_client = MagicMock(spec=HTTPHandler) self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) @@ -2119,6 +2119,28 @@ class TestContextCachingMultiRegionUrls: def setup_method(self): self.caching = ContextCachingEndpoints() + self.context_caching = self.caching + self.mock_logging = MagicMock(spec=Logging) + self.mock_logging.model_call_details = {} + self.mock_client = MagicMock(spec=HTTPHandler) + self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) + self.sample_messages = [ + { + "role": "system", + "content": "You are a helpful assistant", + "cache_control": {"type": "ephemeral"}, + }, + {"role": "user", "content": "Hello, how are you?"}, + ] + self.sample_optional_params = {} + self._token_check_patcher = patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt", + return_value=True, + ) + self._token_check_patcher.start() + + def teardown_method(self): + self._token_check_patcher.stop() @pytest.mark.parametrize("location", ["eu", "us"]) def test_vertex_ai_multi_region_uses_rep_host(self, location): @@ -2161,3 +2183,152 @@ class TestContextCachingMultiRegionUrls: assert url.startswith("https://aiplatform.googleapis.com/") assert "/locations/global/cachedContents" in url + + def test_check_and_create_cache_stashes_creation_metadata(self): + self.mock_logging.model_call_details = {} + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + response = MagicMock() + response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + "usageMetadata": {"totalTokenCount": 10000}, + "createTime": "2025-01-01T00:00:00Z", + "expireTime": "2025-01-02T00:00:00Z", + } + + with ( + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages", + return_value=(cached_messages, non_cached_messages), + ), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj.get_cache_key", + return_value="test_cache_key", + ), + patch.object(self.context_caching, "check_cache", return_value=None), + patch.object( + self.context_caching, + "_get_token_and_url_context_caching", + return_value=("token", "https://test-url.com"), + ), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching", + return_value={"model": "gemini-1.5-pro", "contents": []}, + ), + ): + self.mock_client.post.return_value = response + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=self.sample_optional_params.copy(), + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="token", + ) + + assert self.mock_logging.model_call_details["vertex_ai_cached_content"] == { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + "total_token_count": 10000, + "create_time": "2025-01-01T00:00:00Z", + "expire_time": "2025-01-02T00:00:00Z", + } + + @pytest.mark.asyncio + async def test_async_check_and_create_cache_stashes_creation_metadata(self): + self.mock_logging.model_call_details = {} + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + response = MagicMock() + response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + "usageMetadata": {"totalTokenCount": 10000}, + "createTime": "2025-01-01T00:00:00Z", + "expireTime": "2025-01-02T00:00:00Z", + } + + with ( + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages", + return_value=(cached_messages, non_cached_messages), + ), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj.get_cache_key", + return_value="test_cache_key", + ), + patch.object(self.context_caching, "async_check_cache", return_value=None), + patch.object( + self.context_caching, + "_get_token_and_url_context_caching", + return_value=("token", "https://test-url.com"), + ), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching", + return_value={"model": "gemini-1.5-pro", "contents": []}, + ), + ): + self.mock_async_client.post = AsyncMock(return_value=response) + await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=self.sample_optional_params.copy(), + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="token", + ) + + assert self.mock_logging.model_call_details["vertex_ai_cached_content"] == { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + "total_token_count": 10000, + "create_time": "2025-01-01T00:00:00Z", + "expire_time": "2025-01-02T00:00:00Z", + } + + @pytest.mark.asyncio + async def test_async_check_and_create_cache_reuse_does_not_stash_creation_metadata(self): + self.mock_logging.model_call_details = {} + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + + with ( + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages", + return_value=(cached_messages, non_cached_messages), + ), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj.get_cache_key", + return_value="test_cache_key", + ), + patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"), + ): + await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=self.sample_optional_params.copy(), + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider="vertex_ai", + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="token", + ) + + assert "vertex_ai_cached_content" not in self.mock_logging.model_call_details