From 9dd283534ae4623d32b65283d11931785350e829 Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:17:22 +0900 Subject: [PATCH 01/10] feat(vertex_ai): bill explicit context cache storage per token-hour Vertex AI and Gemini bill explicit context caches for storage by token-hour, on top of the input tokens used to create them. litellm never priced this, so cache storage never reached any spend aggregation. Add a cache_storage_cost_per_token_per_hour pricing key and, when litellm creates a cachedContents resource, record tokens x lifetime (expireTime - createTime, falling back to the requested ttl, then the one-hour provider default) on the request's logging object. completion_cost turns that into additional_costs["cache_storage_cost"] using the deployment or cost-map rate, so it lands in the same request's spend and cost breakdown. Reused caches and models without the key add nothing. --- ci_cd/generate_model_prices_schema.py | 3 + litellm/cost_calculator.py | 30 + litellm/litellm_core_utils/litellm_logging.py | 4 + .../vertex_ai/context_caching/storage_cost.py | 56 ++ .../vertex_ai_context_caching.py | 7 + ...odel_prices_and_context_window_backup.json | 29 +- litellm/types/utils.py | 2 + litellm/utils.py | 1 + model_prices_and_context_window.json | 29 +- model_prices_and_context_window.schema.json | 5 + .../context_caching/test_storage_cost.py | 47 ++ .../test_vertex_ai_context_caching.py | 607 ++++++++---------- tests/unit/test_cost_calculator.py | 125 +++- tests/unit/test_utils.py | 86 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 15 files changed, 645 insertions(+), 390 deletions(-) create mode 100644 litellm/llms/vertex_ai/context_caching/storage_cost.py create mode 100644 tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 8eec07dadda..8530dfd2d08 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -220,6 +220,9 @@ COST_DESCRIPTIONS: dict[str, str] = { ), "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", + "cache_storage_cost_per_token_per_hour": ( + "USD per token per hour an explicitly created context cache (e.g. Vertex AI cachedContents) is stored." + ), "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", } diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b637676836c..e1eeb10ab3d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -299,6 +299,31 @@ def _get_additional_costs( return None +def _context_cache_storage_cost( + litellm_logging_obj: LitellmLoggingObject | None, + model: str, + custom_llm_provider: str | None, + custom_pricing: bool | None, + router_model_id: str | None, +) -> float: + token_hours: Final = litellm_logging_obj.context_cache_storage_token_hours if litellm_logging_obj else 0.0 + if not isinstance(token_hours, int | float) or token_hours <= 0: + return 0.0 + rate: Final = next( + ( + price + for source in ( + _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id), + _cost_map_model_info(model, custom_llm_provider), + ) + if source is not None + and isinstance(price := source.get("cache_storage_cost_per_token_per_hour"), int | float) + ), + 0.0, + ) + return token_hours * rate + + def _transcription_usage_has_token_details( usage_block: Usage | None, ) -> bool: @@ -1847,6 +1872,11 @@ def completion_cost( ) else: additional_costs = None + cache_storage_cost: Final = _context_cache_storage_cost( + litellm_logging_obj, model, custom_llm_provider, custom_pricing, router_model_id + ) + if cache_storage_cost > 0: + additional_costs = {**(additional_costs or {}), "cache_storage_cost": cache_storage_cost} _final_cost = prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar cost_for_built_in_tools = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e15447c8584..766cceceff7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -575,6 +575,10 @@ class Logging(LiteLLMLoggingBaseClass): classifier_input: Mapping[str, JsonValue] | None = None baseline_cache_context: "BaselineCacheContext | None" = None baseline_observation: "CapturedBaselineObservation | None" = None + context_cache_storage_token_hours: float = 0.0 + + def record_context_cache_storage(self, token_hours: float) -> None: + self.context_cache_storage_token_hours = token_hours @property def _defer_async_logging(self) -> bool: diff --git a/litellm/llms/vertex_ai/context_caching/storage_cost.py b/litellm/llms/vertex_ai/context_caching/storage_cost.py new file mode 100644 index 00000000000..7cb74296f94 --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/storage_cost.py @@ -0,0 +1,56 @@ +from datetime import datetime +from typing import Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +DEFAULT_CONTEXT_CACHE_TTL_SECONDS: Final = 3600.0 +_SECONDS_PER_HOUR: Final = 3600.0 + + +class _CachedContentUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + totalTokenCount: int = 0 + + +class _CreatedCachedContent(BaseModel): + model_config = ConfigDict(frozen=True) + + usageMetadata: _CachedContentUsage = _CachedContentUsage() + createTime: datetime | None = None + expireTime: datetime | None = None + + +def _reported_lifetime_seconds(created: _CreatedCachedContent) -> float | None: + if created.createTime is None or created.expireTime is None: + return None + seconds: Final = (created.expireTime - created.createTime).total_seconds() + return seconds if seconds > 0 else None + + +def _requested_ttl_seconds(requested_ttl: str | None) -> float | None: + if requested_ttl is None or not requested_ttl.endswith("s"): + return None + try: + seconds: Final = float(requested_ttl[:-1]) + except ValueError: + return None + return seconds if seconds > 0 else None + + +def context_cache_storage_token_hours(created_cached_content: object, requested_ttl: str | None) -> float: + """Token-hours a newly created cachedContents resource is billed for over its whole TTL. + + The lifetime comes from the create response (expireTime - createTime), then the TTL that was + requested, then the provider default of one hour. + """ + try: + created: Final = _CreatedCachedContent.model_validate(created_cached_content) + except ValidationError: + return 0.0 + lifetime_seconds: Final = ( + _reported_lifetime_seconds(created) + or _requested_ttl_seconds(requested_ttl) + or DEFAULT_CONTEXT_CACHE_TTL_SECONDS + ) + return created.usageMetadata.totalTokenCount * lifetime_seconds / _SECONDS_PER_HOUR 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 2d35dd9b480..7e33bde0c09 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 @@ -21,6 +21,7 @@ from litellm.utils import is_prompt_caching_valid_prompt from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase +from .storage_cost import context_cache_storage_token_hours from .transformation import ( cached_messages_end_on_supported_turn, separate_cached_messages, @@ -425,6 +426,9 @@ class ContextCachingEndpoints(VertexBase): raise VertexAIError(status_code=408, message="Timeout error occurred.") raw_response_cached: Final = response.json() + logging_obj.record_context_cache_storage( + context_cache_storage_token_hours(raw_response_cached, cached_content_request_body.get("ttl")) + ) cached_content_response_obj: Final = VertexAICachedContentResponseObject( name=raw_response_cached.get("name"), model=raw_response_cached.get("model") ) @@ -581,6 +585,9 @@ class ContextCachingEndpoints(VertexBase): raise VertexAIError(status_code=408, message="Timeout error occurred.") raw_response_cached: Final = response.json() + logging_obj.record_context_cache_storage( + context_cache_storage_token_hours(raw_response_cached, cached_content_request_body.get("ttl")) + ) cached_content_response_obj: Final = VertexAICachedContentResponseObject( name=raw_response_cached.get("name"), model=raw_response_cached.get("model") ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8518929a407..463807a68d5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1,5 +1,6 @@ { "sample_spec": { + "cache_storage_cost_per_token_per_hour": 0.0, "code_interpreter_cost_per_session": 0.0, "computer_use_input_cost_per_1k_tokens": 0.0, "computer_use_output_cost_per_1k_tokens": 0.0, @@ -27087,6 +27088,7 @@ "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", @@ -27438,6 +27440,7 @@ }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", @@ -27492,6 +27495,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -27553,6 +27557,7 @@ "cache_read_input_token_cost_batches": 1.5e-08, "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, @@ -27647,6 +27652,7 @@ "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -27935,6 +27941,7 @@ "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, "cache_read_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost_priority": 2.25e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", @@ -28037,6 +28044,7 @@ "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_flex": 2e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28148,6 +28156,7 @@ "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28198,6 +28207,7 @@ "input_cost_per_audio_token_priority": 1.8e-06, "output_cost_per_token_priority": 5.4e-06, "cache_read_input_token_cost_priority": 9e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28270,7 +28280,8 @@ "output_cost_per_token_batches": 4.5e-06, "input_cost_per_token_flex": 7.5e-07, "output_cost_per_token_flex": 4.5e-06, - "cache_read_input_token_cost_flex": 7.5e-08 + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06 }, "vertex_ai/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -28322,6 +28333,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2026-11-19", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -28382,6 +28394,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2027-01-28", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -28442,6 +28455,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28455,6 +28469,7 @@ "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, @@ -28588,6 +28603,7 @@ "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_flex": 2e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -30445,6 +30461,7 @@ "input_cost_per_audio_token_priority": 1.8e-06, "output_cost_per_token_priority": 5.4e-06, "cache_read_input_token_cost_priority": 9e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -30548,7 +30565,8 @@ "output_cost_per_token_batches": 4.5e-06, "input_cost_per_token_flex": 7.5e-07, "output_cost_per_token_flex": 4.5e-06, - "cache_read_input_token_cost_flex": 7.5e-08 + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06 }, "gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -30600,6 +30618,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2026-11-19", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -30660,6 +30679,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2027-01-28", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -30720,6 +30740,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -30733,6 +30754,7 @@ "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, @@ -50860,6 +50882,7 @@ }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", @@ -50914,6 +50937,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -50976,6 +51000,7 @@ "cache_read_input_token_cost_batches": 1.5e-08, "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f8f8388b6f7..791ad9fddcd 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -319,6 +319,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_batches: ReadOnly[float | None] cache_creation_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_creation_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] + cache_storage_cost_per_token_per_hour: ReadOnly[float | None] # explicit context cache storage, e.g. Vertex AI # Smallest prefix this model will actually cache, whatever caching mechanism its provider uses. # Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT. prompt_cache_min_tokens: int | None @@ -3840,6 +3841,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_creation_input_token_cost_batches: float | None = None cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None + cache_storage_cost_per_token_per_hour: float | None = None cache_read_input_audio_token_cost: float | None = None cache_read_input_image_token_cost: float | None = None input_cost_per_character_above_128k_tokens: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 6c0d9064dbe..37a9aa84c01 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6229,6 +6229,7 @@ def get_model_info_helper( cache_creation_input_token_cost_above_1hr=_model_info.get( "cache_creation_input_token_cost_above_1hr", None ), + cache_storage_cost_per_token_per_hour=_model_info.get("cache_storage_cost_per_token_per_hour"), off_peak_pricing=_model_info.get("off_peak_pricing", None), input_cost_per_character=_model_info.get("input_cost_per_character", None), input_cost_per_token_above_128k_tokens=_model_info.get("input_cost_per_token_above_128k_tokens", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8518929a407..463807a68d5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1,5 +1,6 @@ { "sample_spec": { + "cache_storage_cost_per_token_per_hour": 0.0, "code_interpreter_cost_per_session": 0.0, "computer_use_input_cost_per_1k_tokens": 0.0, "computer_use_output_cost_per_1k_tokens": 0.0, @@ -27087,6 +27088,7 @@ "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", @@ -27438,6 +27440,7 @@ }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", @@ -27492,6 +27495,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -27553,6 +27557,7 @@ "cache_read_input_token_cost_batches": 1.5e-08, "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, @@ -27647,6 +27652,7 @@ "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -27935,6 +27941,7 @@ "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, "cache_read_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost_priority": 2.25e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "vertex_ai-language-models", @@ -28037,6 +28044,7 @@ "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_flex": 2e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28148,6 +28156,7 @@ "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28198,6 +28207,7 @@ "input_cost_per_audio_token_priority": 1.8e-06, "output_cost_per_token_priority": 5.4e-06, "cache_read_input_token_cost_priority": 9e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28270,7 +28280,8 @@ "output_cost_per_token_batches": 4.5e-06, "input_cost_per_token_flex": 7.5e-07, "output_cost_per_token_flex": 4.5e-06, - "cache_read_input_token_cost_flex": 7.5e-08 + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06 }, "vertex_ai/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -28322,6 +28333,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2026-11-19", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -28382,6 +28394,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2027-01-28", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -28442,6 +28455,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -28455,6 +28469,7 @@ "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, @@ -28588,6 +28603,7 @@ "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_flex": 2e-07, + "cache_storage_cost_per_token_per_hour": 4.5e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -30445,6 +30461,7 @@ "input_cost_per_audio_token_priority": 1.8e-06, "output_cost_per_token_priority": 5.4e-06, "cache_read_input_token_cost_priority": 9e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -30548,7 +30565,8 @@ "output_cost_per_token_batches": 4.5e-06, "input_cost_per_token_flex": 7.5e-07, "output_cost_per_token_flex": 4.5e-06, - "cache_read_input_token_cost_flex": 7.5e-08 + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06 }, "gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -30600,6 +30618,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2026-11-19", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -30660,6 +30679,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "deprecation_date": "2027-01-28", "search_context_cost_per_query": { "search_context_size_low": 0.014, @@ -30720,6 +30740,7 @@ "input_cost_per_token_priority": 1.35e-06, "output_cost_per_token_priority": 6.75e-06, "cache_read_input_token_cost_priority": 1.35e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "search_context_cost_per_query": { "search_context_size_low": 0.014, "search_context_size_medium": 0.014, @@ -30733,6 +30754,7 @@ "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 1.5e-06, "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, @@ -50860,6 +50882,7 @@ }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-language-models", @@ -50914,6 +50937,7 @@ "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -50976,6 +51000,7 @@ "cache_read_input_token_cost_batches": 1.5e-08, "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5.4e-08, + "cache_storage_cost_per_token_per_hour": 1e-06, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index de3d760d57a..1ed4b588f64 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -256,6 +256,11 @@ "type": "number", "minimum": 0 }, + "cache_storage_cost_per_token_per_hour": { + "type": "number", + "minimum": 0, + "description": "USD per token per hour an explicitly created context cache (e.g. Vertex AI cachedContents) is stored." + }, "citation_cost_per_token": { "type": "number", "minimum": 0 diff --git a/tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py b/tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py new file mode 100644 index 00000000000..b3bfb02b4e1 --- /dev/null +++ b/tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py @@ -0,0 +1,47 @@ +import pytest + +from litellm.llms.vertex_ai.context_caching.storage_cost import ( + context_cache_storage_token_hours, +) + + +@pytest.mark.parametrize( + "created, requested_ttl, expected", + [ + ( + { + "usageMetadata": {"totalTokenCount": 5000}, + "createTime": "2026-01-01T00:00:00.123456789Z", + "expireTime": "2026-01-01T02:00:00.123456789Z", + }, + "60s", + 10000.0, + ), + ({"usageMetadata": {"totalTokenCount": 5000}}, "1800s", 2500.0), + ({"usageMetadata": {"totalTokenCount": 5000}}, "90.5s", 5000 * 90.5 / 3600), + ( + {"usageMetadata": {"totalTokenCount": 4000}, "expireTime": "2026-01-01T02:00:00Z"}, + "900s", + 1000.0, + ), + ({"name": "cachedContents/1"}, "3600s", 0.0), + ({"usageMetadata": {"totalTokenCount": "many"}}, "3600s", 0.0), + ({"usageMetadata": {"totalTokenCount": 5000}, "expireTime": "not-a-time"}, "3600s", 0.0), + ("not-an-object", "3600s", 0.0), + ], +) +def test_token_hours_follow_reported_lifetime_then_requested_ttl(created, requested_ttl, expected): + assert context_cache_storage_token_hours(created, requested_ttl) == pytest.approx(expected) + + +def test_token_hours_default_to_one_hour_without_lifetime_or_ttl(): + # Source: https://cloud.google.com/vertex-ai/generative-ai/docs/context-cache/context-cache-create + # "The default expiration time of a context cache is 60 minutes after it's created" (read 2026-10-07) + assert context_cache_storage_token_hours({"usageMetadata": {"totalTokenCount": 5000}}, None) == 5000.0 + + +@pytest.mark.parametrize("requested_ttl", ["", "s", "abc", "10m", "-5s"]) +def test_unparseable_requested_ttl_falls_back_to_default_lifetime(requested_ttl): + created = {"usageMetadata": {"totalTokenCount": 3600}} + + assert context_cache_storage_token_hours(created, requested_ttl) == context_cache_storage_token_hours(created, None) 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 67d78d6030e..f95b8985198 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 @@ -83,18 +83,10 @@ class TestContextCachingEndpoints: """Teardown for each test method""" self._token_check_patcher.stop() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) - def test_check_and_create_cache_with_cached_content( - self, mock_cache_obj, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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") + def test_check_and_create_cache_with_cached_content(self, mock_cache_obj, mock_separate, custom_llm_provider): """Test check_and_create_cache when cached_content is provided""" # Setup cached_content = "cached_content_123" @@ -129,15 +121,9 @@ class TestContextCachingEndpoints: mock_separate.assert_not_called() mock_cache_obj.get_cache_key.assert_not_called() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) # No cached messages @@ -167,15 +153,9 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache is None - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_existing_cache_found( self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -223,15 +203,9 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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" ) @@ -298,15 +272,9 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert call_args.kwargs["json"]["tools"] == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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.object(ContextCachingEndpoints, "check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") def test_check_and_create_cache_http_error( @@ -331,9 +299,7 @@ class TestContextCachingEndpoints: mock_response = MagicMock() mock_response.status_code = 400 mock_response.text = "Bad Request" - http_error = httpx.HTTPStatusError( - "Error", request=MagicMock(), response=mock_response - ) + http_error = httpx.HTTPStatusError("Error", request=MagicMock(), response=mock_response) self.mock_client.post.side_effect = http_error optional_params = self.sample_optional_params.copy() @@ -361,15 +327,9 @@ class TestContextCachingEndpoints: assert "Bad Request" in str(exc_info.value.message) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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") async def test_async_check_and_create_cache_with_cached_content( self, mock_cache_obj, mock_separate, custom_llm_provider ): @@ -404,15 +364,9 @@ class TestContextCachingEndpoints: assert returned_cache == cached_content @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - async def test_async_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + async def test_async_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test async_check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) @@ -443,15 +397,9 @@ class TestContextCachingEndpoints: assert returned_cache is None @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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.object(ContextCachingEndpoints, "async_check_cache") async def test_async_check_and_create_cache_existing_cache_found( self, mock_async_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -500,23 +448,15 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_create_new_cache( self, mock_get_client, @@ -580,20 +520,12 @@ class TestContextCachingEndpoints: assert call_args.kwargs["json"]["tools"] == self.sample_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_timeout_error( self, mock_get_client, @@ -614,9 +546,7 @@ class TestContextCachingEndpoints: mock_get_token_url.return_value = ("token", "https://test-url.com") # Mock timeout error - self.mock_async_client.post = AsyncMock( - side_effect=httpx.TimeoutException("Timeout") - ) + self.mock_async_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) optional_params = self.sample_optional_params.copy() test_project = "test_project" @@ -642,20 +572,14 @@ class TestContextCachingEndpoints: assert exc_info.value.status_code == 408 assert "Timeout error occurred" in str(exc_info.value.message) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -665,9 +589,7 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): # Execute result = self.context_caching.check_and_create_cache( messages=self.sample_messages, @@ -690,12 +612,8 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -731,12 +649,8 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params in async version when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -772,20 +686,14 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params in async version when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -795,9 +703,7 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the async_check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "async_check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): # Execute result = await self.context_caching.async_check_and_create_cache( messages=self.sample_messages, @@ -820,12 +726,8 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tool_choice_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): """tool_choice is popped from optional_params when cached messages exist.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -837,9 +739,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object( - self.context_caching, "check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): self.context_caching.check_and_create_cache( messages=self.sample_messages, optional_params=optional_params, @@ -857,12 +757,8 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(self, custom_llm_provider): """tool_choice is NOT popped when there are no cached messages (early return).""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -891,12 +787,8 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): """Async equivalent of test_check_and_create_cache_tool_choice_popped_from_optional_params.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -908,9 +800,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object( - self.context_caching, "async_check_cache", return_value="existing_cache" - ): + with 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=optional_params, @@ -929,9 +819,7 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( self, custom_llm_provider ): @@ -962,15 +850,9 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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" ) @@ -1032,15 +914,9 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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" ) @@ -1100,15 +976,9 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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" ) @@ -1160,15 +1030,9 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert "toolConfig" not in call_args.kwargs["json"] - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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" ) @@ -1226,15 +1090,9 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == function_pin - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @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" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @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" ) @@ -1277,9 +1135,7 @@ class TestContextCachingEndpoints: } self.mock_client.post.return_value = mock_response - tool_choice = ToolConfig( - functionCallingConfig=FunctionCallingConfig(mode="ANY") - ) + tool_choice = ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY")) optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = tool_choice @@ -1300,9 +1156,7 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == tool_choice - assert call_args.kwargs["json"]["toolConfig"] == { - "functionCallingConfig": {"mode": "ANY"} - } + assert call_args.kwargs["json"]["toolConfig"] == {"functionCallingConfig": {"mode": "ANY"}} mock_cache_obj.get_cache_key.assert_called_once_with( messages=cached_messages, tools=self.sample_tools, @@ -1310,12 +1164,8 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") @patch.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys( self, @@ -1360,15 +1210,9 @@ class TestContextCachingEndpoints: second_cache_key = check_cache_calls[1].kwargs["cache_key"] assert first_cache_key != second_cache_key - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that context caching is skipped when cached content is below 1024 tokens. Gemini requires a minimum of 1024 tokens for context caching. If the cached @@ -1415,9 +1259,7 @@ class TestContextCachingEndpoints: # Restart the patcher so teardown_method can stop it cleanly self._token_check_patcher.start() - def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum( - self, local_model_cost_map - ): + def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum(self, local_model_cost_map): model = "gemini-2.5-pro" self._token_check_patcher.stop() @@ -1452,16 +1294,10 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider - ): + async def test_async_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that async context caching is skipped when cached content is below 1024 tokens.""" # Stop the default mock so the real token count check runs self._token_check_patcher.stop() @@ -1504,13 +1340,9 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() @pytest.mark.parametrize("is_async", [False, True]) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"]) @pytest.mark.asyncio - async def test_check_and_create_cache_considers_tools_for_min_tokens( - self, custom_llm_provider, is_async - ): + async def test_check_and_create_cache_considers_tools_for_min_tokens(self, custom_llm_provider, is_async): """Test that context caching accounts for tools when validating minimum token count. Fixes #42804: When messages alone are below the threshold, but tools push the total @@ -1535,11 +1367,15 @@ class TestContextCachingEndpoints: "type": "function", "function": { "name": f"synthetic_tool_{i}", - "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " * 8, + "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " + * 8, "parameters": { "type": "object", "properties": { - f"arg_{j}": {"type": "string", "description": "Argument description for caching verification " * 4} + f"arg_{j}": { + "type": "string", + "description": "Argument description for caching verification " * 4, + } for j in range(10) }, "required": [f"arg_{j}" for j in range(5)], @@ -1568,19 +1404,23 @@ class TestContextCachingEndpoints: self.mock_client.post.return_value = mock_response self.mock_async_client.post = AsyncMock(return_value=mock_response) - with patch.object( - self.context_caching, - "_get_token_and_url_context_caching", - return_value=("fake_token", "https://fake.url/cachedContents"), - ), patch.object( - self.context_caching, - "check_cache", - return_value=None, - ), patch.object( - self.context_caching, - "async_check_cache", - new_callable=AsyncMock, - return_value=None, + with ( + patch.object( + self.context_caching, + "_get_token_and_url_context_caching", + return_value=("fake_token", "https://fake.url/cachedContents"), + ), + patch.object( + self.context_caching, + "check_cache", + return_value=None, + ), + patch.object( + self.context_caching, + "async_check_cache", + new_callable=AsyncMock, + return_value=None, + ), ): if is_async: result = await self.context_caching.async_check_and_create_cache( @@ -1624,13 +1464,10 @@ class TestContextCachingEndpoints: post_mock.assert_called_once() call_kwargs = post_mock.call_args.kwargs assert call_kwargs["json"]["tools"] == large_tools - assert call_kwargs["json"]["contents"] == [ - {"role": "user", "parts": [{"text": "Short system instruction."}]} - ] + assert call_kwargs["json"]["contents"] == [{"role": "user", "parts": [{"text": "Short system instruction."}]}] self._token_check_patcher.start() - def _model_turn_final_messages(self, final_cached_role): tool_call = { "id": "call_abc123", @@ -1677,9 +1514,7 @@ class TestContextCachingEndpoints: ] @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) - def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( - self, final_cached_role - ): + def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): """The cachedContents API rejects contents ending on an assistant or tool turn with HTTP 400 "Requests ending with a model turn are not supported", so the request must proceed uncached instead of failing. @@ -1712,9 +1547,7 @@ class TestContextCachingEndpoints: @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( - self, final_cached_role - ): + async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): """Async variant: an unsupported terminal turn skips caching instead of failing.""" all_messages = self._model_turn_final_messages(final_cached_role) optional_params = self.sample_optional_params.copy() @@ -1772,14 +1605,8 @@ def test_cached_messages_end_on_supported_turn(): ) is True ) - assert ( - cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) - is False - ) - assert ( - cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) - is False - ) + assert cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) is False + assert cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) is False assert cached_messages_end_on_supported_turn([]) is False @@ -1793,13 +1620,9 @@ class TestCheckCachePagination: self.mock_client = MagicMock(spec=HTTPHandler) self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_finds_cache_on_second_page( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_finds_cache_on_second_page(self, mock_get_token_url, custom_llm_provider): """Test that check_cache correctly handles pagination and finds cache on second page""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1848,13 +1671,9 @@ class TestCheckCachePagination: second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"] assert "pageToken=token_page_2" in second_call_url - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_stops_when_no_next_token( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): """Test that check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1889,13 +1708,9 @@ class TestCheckCachePagination: assert result is None assert self.mock_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_multiple_pages( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_multiple_pages(self, mock_get_token_url, custom_llm_provider): """Test that check_cache correctly iterates through multiple pages""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1940,9 +1755,7 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 3 @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") async def test_async_check_cache_pagination_finds_cache_on_second_page( self, mock_get_token_url, custom_llm_provider @@ -1972,9 +1785,7 @@ class TestCheckCachePagination: } # Setup mock async client to return different responses - self.mock_async_client.get = AsyncMock( - side_effect=[first_page_response, second_page_response] - ) + self.mock_async_client.get = AsyncMock(side_effect=[first_page_response, second_page_response]) # Execute result = await self.context_caching.async_check_cache( @@ -1998,13 +1809,9 @@ class TestCheckCachePagination: assert "pageToken=token_page_2" in second_call_url @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_stops_when_no_next_token( - self, mock_get_token_url, custom_llm_provider - ): + async def test_async_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): """Test that async_check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -2039,29 +1846,21 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): + def test_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_client.get.side_effect = [ - create_page_response(i) for i in range(100) - ] + self.mock_client.get.side_effect = [create_page_response(i) for i in range(100)] result = self.context_caching.check_cache( cache_key=cache_key_to_find, @@ -2080,29 +1879,21 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 100 @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_max_pages_limit( - self, mock_get_token_url, custom_llm_provider - ): + async def test_async_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_async_client.get = AsyncMock( - side_effect=[create_page_response(i) for i in range(100)] - ) + self.mock_async_client.get = AsyncMock(side_effect=[create_page_response(i) for i in range(100)]) result = await self.context_caching.async_check_cache( cache_key=cache_key_to_find, @@ -2121,8 +1912,6 @@ class TestCheckCachePagination: assert self.mock_async_client.get.call_count == 100 - - class TestVertexAIGlobalLocation: """Test global location handling in context caching.""" @@ -2148,9 +1937,7 @@ class TestVertexAIGlobalLocation: # Assert correct URL format for global expected_url = "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/cachedContents" assert url == expected_url, f"Expected {expected_url}, got {url}" - assert ( - "global-aiplatform" not in url - ), "URL should not contain 'global-aiplatform' prefix" + assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" def test_regional_location_url_construction_v1(self): """Test that regional location uses correct URL (with location prefix) for v1 API.""" @@ -2193,11 +1980,11 @@ class TestVertexAIGlobalLocation: ) # Assert correct URL format for global with beta API - expected_url = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" + expected_url = ( + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" + ) assert url == expected_url, f"Expected {expected_url}, got {url}" - assert ( - "global-aiplatform" not in url - ), "URL should not contain 'global-aiplatform' prefix" + assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" def test_gemini_context_caching_with_custom_api_base_passes_model(self): """Gemini context caching with custom api_base must pass model to _check_custom_proxy. @@ -2289,3 +2076,127 @@ class TestContextCachingMultiRegionUrls: assert url.startswith("https://aiplatform.googleapis.com/") assert "/locations/global/cachedContents" in url + + +_STORAGE_TEST_MODEL: Final = "gemini-2.5-pro" +_STORAGE_TEST_CACHE_NAME: Final = "cachedContents/storage-cost" + + +def _storage_logging_obj() -> Logging: + logging_obj: Final = Logging( + model=_STORAGE_TEST_MODEL, + messages=[], + stream=False, + call_type="acompletion", + start_time=None, + litellm_call_id="context-cache-storage", + function_id="context-cache-storage", + ) + logging_obj.update_environment_variables( + model=_STORAGE_TEST_MODEL, litellm_params={}, optional_params={}, custom_llm_provider="gemini" + ) + return logging_obj + + +def _storage_messages() -> list: + return [ + { + "role": "system", + "content": [{"type": "text", "text": "policy " * 6000, "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "Hello"}, + ] + + +def _json_response(method: str, payload: dict) -> httpx.Response: + return httpx.Response(200, json=payload, request=httpx.Request(method, "https://example.invalid")) + + +_CREATED_CACHE: Final = { + "name": _STORAGE_TEST_CACHE_NAME, + "model": f"models/{_STORAGE_TEST_MODEL}", + "usageMetadata": {"totalTokenCount": 7200}, + "createTime": "2026-01-01T00:00:00Z", + "expireTime": "2026-01-01T00:30:00Z", +} + + +def _create_kwargs(logging_obj: Logging, client: object) -> dict: + return { + "messages": _storage_messages(), + "optional_params": {}, + "api_key": "test_key", + "api_base": None, + "model": _STORAGE_TEST_MODEL, + "client": client, + "timeout": 30.0, + "logging_obj": logging_obj, + "custom_llm_provider": "gemini", + "vertex_project": None, + "vertex_location": None, + "vertex_auth_header": None, + } + + +def test_creating_a_cache_records_its_storage_token_hours(local_model_cost_map): + logging_obj: Final = _storage_logging_obj() + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _json_response("GET", {}) + client.post.return_value = _json_response("POST", _CREATED_CACHE) + + _, _, cache_name = ContextCachingEndpoints().check_and_create_cache(**_create_kwargs(logging_obj, client)) + + assert cache_name == _STORAGE_TEST_CACHE_NAME + assert logging_obj.context_cache_storage_token_hours == pytest.approx(7200 * 0.5) + + +@pytest.mark.asyncio +async def test_async_creating_a_cache_records_its_storage_token_hours(local_model_cost_map): + logging_obj: Final = _storage_logging_obj() + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_json_response("GET", {})) + client.post = AsyncMock(return_value=_json_response("POST", _CREATED_CACHE)) + + _, _, cache_name = await ContextCachingEndpoints().async_check_and_create_cache( + **_create_kwargs(logging_obj, client) + ) + + assert cache_name == _STORAGE_TEST_CACHE_NAME + assert logging_obj.context_cache_storage_token_hours == pytest.approx(7200 * 0.5) + + +def _listing_with(created_body: dict) -> dict: + return {"cachedContents": [{"name": _STORAGE_TEST_CACHE_NAME, "displayName": created_body["displayName"]}]} + + +def test_reusing_an_existing_cache_records_no_storage(local_model_cost_map): + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _json_response("GET", {}) + client.post.return_value = _json_response("POST", _CREATED_CACHE) + ContextCachingEndpoints().check_and_create_cache(**_create_kwargs(_storage_logging_obj(), client)) + client.get.return_value = _json_response("GET", _listing_with(client.post.call_args.kwargs["json"])) + reusing_logging_obj: Final = _storage_logging_obj() + + _, _, cache_name = ContextCachingEndpoints().check_and_create_cache(**_create_kwargs(reusing_logging_obj, client)) + + assert cache_name == _STORAGE_TEST_CACHE_NAME + assert client.post.call_count == 1 + assert reusing_logging_obj.context_cache_storage_token_hours == 0.0 + + +@pytest.mark.asyncio +async def test_async_reusing_an_existing_cache_records_no_storage(local_model_cost_map): + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_json_response("GET", {})) + client.post = AsyncMock(return_value=_json_response("POST", _CREATED_CACHE)) + await ContextCachingEndpoints().async_check_and_create_cache(**_create_kwargs(_storage_logging_obj(), client)) + client.get.return_value = _json_response("GET", _listing_with(client.post.call_args.kwargs["json"])) + reusing_logging_obj: Final = _storage_logging_obj() + + _, _, cache_name = await ContextCachingEndpoints().async_check_and_create_cache( + **_create_kwargs(reusing_logging_obj, client) + ) + + assert cache_name == _STORAGE_TEST_CACHE_NAME + assert client.post.call_count == 1 + assert reusing_logging_obj.context_cache_storage_token_hours == 0.0 diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 9e348c0579a..726fb4c33c3 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -189,7 +189,9 @@ def test_response_cost_calculator_keeps_optional_params_out_of_hidden_params(): assert optional_params["aws_session_token"] == "session-secret" -def test_embedding_success_logging_and_spend_log_carry_no_forwarded_credentials(monkeypatch: pytest.MonkeyPatch) -> None: +def test_embedding_success_logging_and_spend_log_carry_no_forwarded_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: from litellm.proxy import proxy_server from litellm.proxy.spend_tracking.spend_tracking_utils import _get_proxy_server_request_for_spend_logs_payload @@ -238,10 +240,6 @@ def test_embedding_success_logging_and_spend_log_carry_no_forwarded_credentials( assert logging_obj.optional_params["extra_headers"] == {"x-goog-api-key": "goog-secret"} - - - - def test_realtime_stream_combines_text_and_audio_token_details(): """Realtime response.done usage with input_token_details / output_token_details.""" from litellm.cost_calculator import RealtimeAPITokenUsageProcessor @@ -1356,8 +1354,6 @@ def test_bedrock_cost_calculator_comparison_with_without_cache(): print(f"Cost with cache: {cost_with_cache}") - - def test_gemini_25_explicit_caching_cost_direct_usage(): """ Test that Gemini 2.5 models correctly calculate costs with explicit caching. @@ -1994,8 +1990,6 @@ def test_cost_margin_with_discount(monkeypatch): print(f" - Expected: ${expected_cost:.6f}") - - def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_map): """Test that completion_cost extracts service_tier from completion_response object.""" from litellm import completion_cost @@ -2745,8 +2739,6 @@ def test_gemini_without_cache_tokens_details(): print("✅ Gemini without cacheTokensDetails works correctly") - - def test_additional_costs_only_for_azure_ai(_local_model_cost_map): """ Test that _get_additional_costs is only called for azure_ai provider. @@ -3291,9 +3283,7 @@ def test_cost_per_token_resolves_per_second_rate_precedence( model: Final = "test-chat-per-second-rate-precedence" entry: Final = {**pricing_fields, "litellm_provider": "together_ai", "mode": "chat"} - litellm.register_model( - model_cost={model: entry} - ) + litellm.register_model(model_cost={model: entry}) assert cost_per_token( model=model, @@ -4061,7 +4051,10 @@ def test_completion_cost_region_without_its_own_row_prices_mantle_claude_from_th custom_llm_provider="bedrock_mantle", region_name="us-east-1", ) == pytest.approx(expected), deployment - assert litellm.get_model_info(f"bedrock_mantle/us-east-1/{model}", "bedrock_mantle")["key"] == f"bedrock_mantle/{model}" + assert ( + litellm.get_model_info(f"bedrock_mantle/us-east-1/{model}", "bedrock_mantle")["key"] + == f"bedrock_mantle/{model}" + ) @pytest.mark.parametrize("model", ["anthropic.claude-opus-5-5", "anthropic.claude-sonnet-5-5"]) @@ -4985,9 +4978,7 @@ def test_xai_batch_tier_discounts_the_long_context_rate_like_the_flat_batch_rate assert info[f"{prefix}_above_200k_tokens_batches"] < info[f"{prefix}_above_200k_tokens"] -@pytest.mark.parametrize( - ("prompt_tokens", "tier"), [(200_000, "_above_200k_tokens_batches"), (199_999, "_batches")] -) +@pytest.mark.parametrize(("prompt_tokens", "tier"), [(200_000, "_above_200k_tokens_batches"), (199_999, "_batches")]) def test_xai_batch_cost_calculator_bills_the_200k_batch_tier_inclusively( _local_model_cost_map: None, prompt_tokens: int, tier: str ) -> None: @@ -5839,3 +5830,101 @@ def test_completion_cost_bills_base_when_gemini_serves_on_demand( ) assert cost == pytest.approx(100 * 0.001 + 50 * 0.002) + + +_STORAGE_RATE_KEY: Final = "cache_storage_cost_per_token_per_hour" + + +@pytest.fixture +def _storage_cost_map(_local_model_cost_map: None): + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def _first_chat_model(provider_prefix: str, *, with_storage_rate: bool) -> str: + return next( + name + for name, entry in litellm.model_cost.items() + if _is_priced_chat_entry(entry, provider_prefix) and (_STORAGE_RATE_KEY in entry) == with_storage_rate + ) + + +def _is_priced_chat_entry(entry: object, provider_prefix: str) -> bool: + return ( + isinstance(entry, dict) + and entry.get("mode") == "chat" + and str(entry.get("litellm_provider", "")).startswith(provider_prefix) + and entry.get("input_cost_per_token") is not None + ) + + +def _cache_storage_logging_obj(model: str, token_hours: float) -> Logging: + logging_obj: Final = Logging( + model=model, + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="context-cache-storage", + function_id="f", + ) + if token_hours: + logging_obj.record_context_cache_storage(token_hours) + return logging_obj + + +def _chat_cost(model: str, logging_obj: Logging, **kwargs: object) -> float: + usage: Final = Usage(prompt_tokens=1_000, completion_tokens=100, total_tokens=1_100) + return completion_cost( + completion_response=ModelResponse(model=model, usage=usage), + model=model, + litellm_logging_obj=logging_obj, + **kwargs, + ) + + +def test_context_cache_storage_is_billed_on_the_creating_request(_storage_cost_map: None) -> None: + model: Final = _first_chat_model("vertex_ai", with_storage_rate=True) + rate: Final = litellm.get_model_info(model)[_STORAGE_RATE_KEY] + creating_request: Final = _cache_storage_logging_obj(model, token_hours=3_600.0) + + with_storage: Final = _chat_cost(model, creating_request) + without_storage: Final = _chat_cost(model, _cache_storage_logging_obj(model, token_hours=0.0)) + + assert rate is not None and rate > 0 + assert with_storage - without_storage == pytest.approx(3_600.0 * rate) + assert creating_request.cost_breakdown["additional_costs"] == pytest.approx({"cache_storage_cost": 3_600.0 * rate}) + assert creating_request.cost_breakdown["total_cost"] == pytest.approx(with_storage) + + +def test_context_cache_storage_costs_nothing_for_models_without_a_storage_rate(_storage_cost_map: None) -> None: + model: Final = _first_chat_model("gemini", with_storage_rate=False) + creating_request: Final = _cache_storage_logging_obj(model, token_hours=3_600.0) + + with_storage: Final = _chat_cost(model, creating_request) + + assert with_storage == _chat_cost(model, _cache_storage_logging_obj(model, token_hours=0.0)) + assert creating_request.cost_breakdown.get("additional_costs") is None + + +def test_context_cache_storage_uses_the_deployment_storage_rate(_storage_cost_map: None) -> None: + model: Final = _first_chat_model("vertex_ai", with_storage_rate=True) + deployment_rate: Final = litellm.get_model_info(model)[_STORAGE_RATE_KEY] * 3 + litellm.register_model( + { + "storage-cost-deployment": { + "litellm_provider": "vertex_ai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + _STORAGE_RATE_KEY: deployment_rate, + } + } + ) + pricing: Final = {"custom_pricing": True, "router_model_id": "storage-cost-deployment"} + + with_storage: Final = _chat_cost(model, _cache_storage_logging_obj(model, token_hours=500.0), **pricing) + without_storage: Final = _chat_cost(model, _cache_storage_logging_obj(model, token_hours=0.0), **pricing) + + assert with_storage - without_storage == pytest.approx(500.0 * deployment_rate) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index c899083ca80..cb99e178dc4 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -63,7 +63,7 @@ from litellm.types.utils import ( bedrock_batch_litellm_params, ) from litellm.types.videos.main import VideoObject -from litellm.utils import( +from litellm.utils import ( _invalidate_model_cost_lowercase_map, CustomStreamWrapper, filter_out_litellm_params, @@ -790,6 +790,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_creation_input_token_cost_flex": {"type": "number"}, "cache_creation_input_token_cost_priority": {"type": "number"}, "cache_creation_input_token_cost_ultrafast": {"type": "number"}, + "cache_storage_cost_per_token_per_hour": {"type": "number"}, "cache_read_input_token_cost": {"type": "number"}, "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, @@ -3164,9 +3165,7 @@ class TestExtraBodyCannotOverrideModel: } ] - untouched = litellm.get_optional_params( - model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools() - ) + untouched = litellm.get_optional_params(model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools()) assert untouched["tools"][0]["function"]["custom_marker"] == "LEAK", untouched result = litellm.get_optional_params( @@ -4157,7 +4156,9 @@ def test_is_prompt_caching_valid_prompt_stops_counting_once_the_minimum_is_reach assert is_prompt_caching_valid_prompt(model="claude-opus-4-8", messages=long_prompt, min_token_count=1024) is True assert sum(counted_messages) < len(long_prompt), sum(counted_messages) - full_count = litellm.token_counter(model="claude-opus-4-8", messages=long_prompt, use_default_image_token_count=True) + full_count = litellm.token_counter( + model="claude-opus-4-8", messages=long_prompt, use_default_image_token_count=True + ) assert ( is_prompt_caching_valid_prompt(model="claude-opus-4-8", messages=long_prompt, min_token_count=full_count) is True @@ -4620,7 +4621,13 @@ _SUCCESS_RESPONSES_BY_CALL_TYPE: Final = ( pytest.param(EmbeddingResponse(model="text-embedding-3-small"), CallTypes.aembedding, id="embedding"), pytest.param( ResponsesAPIResponse( - id="resp_abc", created_at=1, output=[], parallel_tool_calls=False, tool_choice="auto", tools=[], model="gpt-5.6" + id="resp_abc", + created_at=1, + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + model="gpt-5.6", ), CallTypes.aresponses, id="responses", @@ -4648,7 +4655,9 @@ async def test_success_deployment_hook_raising_keeps_response_and_runs_later_hoo assert result is response assert second_hook.seen_responses == (response,) - failure_logs: Final = tuple(r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message) + failure_logs: Final = tuple( + r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message + ) assert len(failure_logs) == 1 assert "_ChatShapedSuccessDeploymentHook" in failure_logs[0].message assert str(call_type) in failure_logs[0].message @@ -4785,17 +4794,13 @@ async def test_wrapper_async_logs_converted_responses_stream_with_standard_loggi ({"metadata": None}, True, False), ], ) -def test_is_litellm_router_call_is_async_aware( - kwargs: Mapping[str, object], is_async: bool, expected: bool -) -> None: +def test_is_litellm_router_call_is_async_aware(kwargs: Mapping[str, object], is_async: bool, expected: bool) -> None: assert _is_litellm_router_call(kwargs, is_async=is_async) is expected @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) -async def test_router_aresponses_does_not_run_sdk_retries( - monkeypatch: pytest.MonkeyPatch, stream: bool -) -> None: +async def test_router_aresponses_does_not_run_sdk_retries(monkeypatch: pytest.MonkeyPatch, stream: bool) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() model_list: Final = [ @@ -4809,9 +4814,7 @@ async def test_router_aresponses_does_not_run_sdk_retries( }, } ] - router: Final = litellm.Router( - model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True - ) + router: Final = litellm.Router(model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True) try: with respx.mock(assert_all_called=True) as respx_mock: @@ -4847,9 +4850,7 @@ def test_router_responses_keeps_sdk_retries_for_sync_router_call( }, } ] - router: Final = litellm.Router( - model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True - ) + router: Final = litellm.Router(model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True) try: with respx.mock(assert_all_called=True) as respx_mock: @@ -6693,7 +6694,9 @@ def test_function_setup_logs_the_search_query_edit_prompt_and_ocr_document_summa @pytest.mark.parametrize("original_function", ("atext_completion", "text_completion")) -def test_function_setup_without_a_prompt_leaves_the_missing_prompt_to_request_validation(original_function: str) -> None: +def test_function_setup_without_a_prompt_leaves_the_missing_prompt_to_request_validation( + original_function: str, +) -> None: assert _logged_request_messages(original_function, model="gpt-4o") is None @@ -6735,6 +6738,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown(): """ @@ -6748,14 +6752,17 @@ def setup_and_teardown(): loop.close() asyncio.set_event_loop(None) + MODEL: Final = "anthropic/claude-haiku-4-5" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_none(): """Test that None is returned as-is.""" result = validate_chat_completion_tool_choice(None, model=MODEL) assert result is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_string(): """Test that string values are returned as-is.""" @@ -6763,6 +6770,7 @@ def test_validate_tool_choice_string(): assert validate_chat_completion_tool_choice("none", model=MODEL) == "none" assert validate_chat_completion_tool_choice("required", model=MODEL) == "required" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_standard_dict(): """Test standard OpenAI format with function.""" @@ -6770,6 +6778,7 @@ def test_validate_tool_choice_standard_dict(): result = validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert result == tool_choice + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_cursor_format(): """Cursor IDE format {"type": "auto"} is unwrapped to the bare string.""" @@ -6777,6 +6786,7 @@ def test_validate_tool_choice_cursor_format(): assert validate_chat_completion_tool_choice({"type": "none"}, model=MODEL) == "none" assert validate_chat_completion_tool_choice({"type": "required"}, model=MODEL) == "required" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize( "tool_choice", @@ -6798,6 +6808,7 @@ def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): assert exc_info.value.status_code == 400 assert exc_info.value.model == MODEL + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("tool_choice", [123, []]) def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): @@ -6808,6 +6819,7 @@ def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert exc_info.value.status_code == 400 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_without_model_is_still_a_400(): """Callers that predate the model argument keep getting a 400, with an empty model on the error.""" @@ -6816,12 +6828,14 @@ def test_validate_tool_choice_without_model_is_still_a_400(): assert exc_info.value.status_code == 400 assert exc_info.value.model == "" + @pytest.fixture() def _vcr_outcome_gate_local_testing(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -6874,6 +6888,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -6894,6 +6909,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown_local_testing(): """ @@ -6916,6 +6932,7 @@ def setup_and_teardown_local_testing(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -6941,6 +6958,7 @@ def test_vertex_projects(): assert "vertex_ai_project" in optional_params assert "vertex_ai_location" in optional_params + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -6956,6 +6974,7 @@ def test_bedrock_embed_v2_regular(): print(f"received optional_params: {optional_params}") assert optional_params == {"dimensions": 512} + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -6974,6 +6993,7 @@ def test_bedrock_embed_v2_with_drop_params(): print(f"received optional_params: {optional_params}") assert optional_params == {"dimensions": 512, "embeddingTypes": ["binary"]} + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -6995,6 +7015,7 @@ def test_openai_non_text_embedding_3_with_allowed_openai_params(): print(f"received optional_params: {optional_params}") assert optional_params.get("dimensions") == 1024 + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7021,6 +7042,7 @@ def test_openai_non_text_embedding_3_without_allowed_openai_params_raises(): finally: litellm.drop_params = prev_drop_params + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7049,6 +7071,7 @@ def test_openai_non_text_embedding_3_drop_params_per_call(): finally: litellm.drop_params = prev_drop_params + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7076,12 +7099,14 @@ def test_openai_non_text_embedding_3_drop_params_global(): finally: litellm.drop_params = prev_drop_params + @pytest.fixture() def _vcr_outcome_gate_search_tests(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.mark.usefixtures("_vcr_outcome_gate_search_tests") def test_search_tool_name_in_all_litellm_params(): """ @@ -7091,6 +7116,7 @@ def test_search_tool_name_in_all_litellm_params(): """ assert "search_tool_name" in all_litellm_params + @pytest.mark.usefixtures("_vcr_outcome_gate_search_tests") def test_filter_out_search_tool_name(): """ @@ -7117,6 +7143,7 @@ def test_filter_out_search_tool_name(): assert filtered["query"] == "latest ai developments" assert filtered["max_results"] == 5 + @pytest.mark.asyncio async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None: """Chat over the Responses bridge exits two @client wrappers with one logging object. Issue @@ -7179,3 +7206,22 @@ async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: assert len(counting_logger.logged_results) == 1, counting_logger.logged_results assert counting_logger.logged_results[0] is inner_result + + +def test_get_model_info_surfaces_context_cache_storage_rate(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + priced: Final = { + name: entry["cache_storage_cost_per_token_per_hour"] + for name, entry in litellm.model_cost.items() + if isinstance(entry, dict) and "cache_storage_cost_per_token_per_hour" in entry + } + + surfaced: Final = { + name: litellm.get_model_info(name).get("cache_storage_cost_per_token_per_hour") for name in priced + } + litellm.get_model_info.cache_clear() + + assert priced + assert surfaced == priced diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cbaea37d409..7599aa31f52 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35557,6 +35557,8 @@ export interface components { cache_read_input_token_cost_priority?: number | null; /** Cache Read Input Token Cost Ultrafast */ cache_read_input_token_cost_ultrafast?: number | null; + /** Cache Storage Cost Per Token Per Hour */ + cache_storage_cost_per_token_per_hour?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Client Id */ @@ -50454,6 +50456,8 @@ export interface components { cache_read_input_token_cost_priority?: number | null; /** Cache Read Input Token Cost Ultrafast */ cache_read_input_token_cost_ultrafast?: number | null; + /** Cache Storage Cost Per Token Per Hour */ + cache_storage_cost_per_token_per_hour?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Client Id */ From db5c264591c8e81e35f545d98f0f0a21301faa79 Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:31:14 +0900 Subject: [PATCH 02/10] fix(cost_map): keep schema in sync with base generator and limit test diffs to additions cost-map-guard runs the base branch's generator under pull_request_target, so a new COST_DESCRIPTIONS entry in this PR made the committed schema diverge from what CI renders. Drop the description and regenerate the schema. Also revert unrelated reformatting of existing test files so the diff only adds the new tests. --- ci_cd/generate_model_prices_schema.py | 3 - model_prices_and_context_window.schema.json | 3 +- .../test_vertex_ai_context_caching.py | 485 +++++++++++++----- tests/unit/test_cost_calculator.py | 27 +- tests/unit/test_utils.py | 66 +-- 5 files changed, 389 insertions(+), 195 deletions(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 8530dfd2d08..8eec07dadda 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -220,9 +220,6 @@ COST_DESCRIPTIONS: dict[str, str] = { ), "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", - "cache_storage_cost_per_token_per_hour": ( - "USD per token per hour an explicitly created context cache (e.g. Vertex AI cachedContents) is stored." - ), "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 1ed4b588f64..58c7e4928c6 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -258,8 +258,7 @@ }, "cache_storage_cost_per_token_per_hour": { "type": "number", - "minimum": 0, - "description": "USD per token per hour an explicitly created context cache (e.g. Vertex AI cachedContents) is stored." + "minimum": 0 }, "citation_cost_per_token": { "type": "number", 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 f95b8985198..77907e0ca19 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 @@ -83,10 +83,18 @@ class TestContextCachingEndpoints: """Teardown for each test method""" self._token_check_patcher.stop() - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") - def test_check_and_create_cache_with_cached_content(self, mock_cache_obj, mock_separate, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" + ) + def test_check_and_create_cache_with_cached_content( + self, mock_cache_obj, mock_separate, custom_llm_provider + ): """Test check_and_create_cache when cached_content is provided""" # Setup cached_content = "cached_content_123" @@ -121,9 +129,15 @@ class TestContextCachingEndpoints: mock_separate.assert_not_called() mock_cache_obj.get_cache_key.assert_not_called() - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") - def test_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + def test_check_and_create_cache_no_cached_messages( + self, mock_separate, custom_llm_provider + ): """Test check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) # No cached messages @@ -153,9 +167,15 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache is None - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_existing_cache_found( self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -203,9 +223,15 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" ) @@ -272,9 +298,15 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert call_args.kwargs["json"]["tools"] == self.sample_tools - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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.object(ContextCachingEndpoints, "check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") def test_check_and_create_cache_http_error( @@ -299,7 +331,9 @@ class TestContextCachingEndpoints: mock_response = MagicMock() mock_response.status_code = 400 mock_response.text = "Bad Request" - http_error = httpx.HTTPStatusError("Error", request=MagicMock(), response=mock_response) + http_error = httpx.HTTPStatusError( + "Error", request=MagicMock(), response=mock_response + ) self.mock_client.post.side_effect = http_error optional_params = self.sample_optional_params.copy() @@ -327,9 +361,15 @@ class TestContextCachingEndpoints: assert "Bad Request" in str(exc_info.value.message) @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" + ) async def test_async_check_and_create_cache_with_cached_content( self, mock_cache_obj, mock_separate, custom_llm_provider ): @@ -364,9 +404,15 @@ class TestContextCachingEndpoints: assert returned_cache == cached_content @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") - async def test_async_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + async def test_async_check_and_create_cache_no_cached_messages( + self, mock_separate, custom_llm_provider + ): """Test async_check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) @@ -397,9 +443,15 @@ class TestContextCachingEndpoints: assert returned_cache is None @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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.object(ContextCachingEndpoints, "async_check_cache") async def test_async_check_and_create_cache_existing_cache_found( self, mock_async_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -448,15 +500,23 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" + ) async def test_async_check_and_create_cache_create_new_cache( self, mock_get_client, @@ -520,12 +580,20 @@ class TestContextCachingEndpoints: assert call_args.kwargs["json"]["tools"] == self.sample_tools @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" + ) async def test_async_check_and_create_cache_timeout_error( self, mock_get_client, @@ -546,7 +614,9 @@ class TestContextCachingEndpoints: mock_get_token_url.return_value = ("token", "https://test-url.com") # Mock timeout error - self.mock_async_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) + self.mock_async_client.post = AsyncMock( + side_effect=httpx.TimeoutException("Timeout") + ) optional_params = self.sample_optional_params.copy() test_project = "test_project" @@ -572,14 +642,20 @@ class TestContextCachingEndpoints: assert exc_info.value.status_code == 408 assert "Timeout error occurred" in str(exc_info.value.message) - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - def test_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + def test_check_and_create_cache_tools_popped_from_optional_params( + self, custom_llm_provider + ): """Test that tools are properly popped from optional_params when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [self.sample_messages[0]] # System message with cache_control + cached_messages = [ + self.sample_messages[0] + ] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -589,7 +665,9 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the check_cache to return existing cache so we don't make HTTP calls - with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): + with patch.object( + self.context_caching, "check_cache", return_value="existing_cache" + ): # Execute result = self.context_caching.check_and_create_cache( messages=self.sample_messages, @@ -612,8 +690,12 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - def test_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + def test_check_and_create_cache_tools_not_popped_when_no_cached_messages( + self, custom_llm_provider + ): """Test that tools are NOT popped from optional_params when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -649,8 +731,12 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages( + self, custom_llm_provider + ): """Test that tools are NOT popped from optional_params in async version when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -686,14 +772,20 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - async def test_async_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + async def test_async_check_and_create_cache_tools_popped_from_optional_params( + self, custom_llm_provider + ): """Test that tools are properly popped from optional_params in async version when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [self.sample_messages[0]] # System message with cache_control + cached_messages = [ + self.sample_messages[0] + ] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -703,7 +795,9 @@ class TestContextCachingEndpoints: test_location = "us-central1" # Mock the async_check_cache to return existing cache so we don't make HTTP calls - with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): + with patch.object( + self.context_caching, "async_check_cache", return_value="existing_cache" + ): # Execute result = await self.context_caching.async_check_and_create_cache( messages=self.sample_messages, @@ -726,8 +820,12 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - def test_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + def test_check_and_create_cache_tool_choice_popped_from_optional_params( + self, custom_llm_provider + ): """tool_choice is popped from optional_params when cached messages exist.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -739,7 +837,9 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): + with patch.object( + self.context_caching, "check_cache", return_value="existing_cache" + ): self.context_caching.check_and_create_cache( messages=self.sample_messages, optional_params=optional_params, @@ -757,8 +857,12 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(self, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( + self, custom_llm_provider + ): """tool_choice is NOT popped when there are no cached messages (early return).""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -787,8 +891,12 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params( + self, custom_llm_provider + ): """Async equivalent of test_check_and_create_cache_tool_choice_popped_from_optional_params.""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -800,7 +908,9 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = {"functionCallingConfig": {"mode": "ANY"}} - with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): + with 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=optional_params, @@ -819,7 +929,9 @@ class TestContextCachingEndpoints: assert "tool_choice" not in optional_params @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( self, custom_llm_provider ): @@ -850,9 +962,15 @@ class TestContextCachingEndpoints: assert optional_params.get("tool_choice") == tool_choice - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" ) @@ -914,9 +1032,15 @@ class TestContextCachingEndpoints: ) @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" ) @@ -976,9 +1100,15 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" ) @@ -1030,9 +1160,15 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert "toolConfig" not in call_args.kwargs["json"] - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" ) @@ -1090,9 +1226,15 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == function_pin - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @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") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @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" ) @@ -1135,7 +1277,9 @@ class TestContextCachingEndpoints: } self.mock_client.post.return_value = mock_response - tool_choice = ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY")) + tool_choice = ToolConfig( + functionCallingConfig=FunctionCallingConfig(mode="ANY") + ) optional_params = self.sample_optional_params.copy() optional_params["tool_choice"] = tool_choice @@ -1156,7 +1300,9 @@ class TestContextCachingEndpoints: call_args = self.mock_client.post.call_args assert call_args.kwargs["json"]["toolConfig"] == tool_choice - assert call_args.kwargs["json"]["toolConfig"] == {"functionCallingConfig": {"mode": "ANY"}} + assert call_args.kwargs["json"]["toolConfig"] == { + "functionCallingConfig": {"mode": "ANY"} + } mock_cache_obj.get_cache_key.assert_called_once_with( messages=cached_messages, tools=self.sample_tools, @@ -1164,8 +1310,12 @@ class TestContextCachingEndpoints: model="gemini-1.5-pro", ) - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) @patch.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_distinct_tool_choices_use_distinct_keys( self, @@ -1210,9 +1360,15 @@ class TestContextCachingEndpoints: second_cache_key = check_cache_calls[1].kwargs["cache_key"] assert first_cache_key != second_cache_key - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") - def test_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + def test_check_and_create_cache_skips_when_below_min_tokens( + self, mock_separate, custom_llm_provider + ): """Test that context caching is skipped when cached content is below 1024 tokens. Gemini requires a minimum of 1024 tokens for context caching. If the cached @@ -1259,7 +1415,9 @@ class TestContextCachingEndpoints: # Restart the patcher so teardown_method can stop it cleanly self._token_check_patcher.start() - def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum(self, local_model_cost_map): + def test_check_and_create_cache_skips_between_default_and_gemini_2_5_pro_minimum( + self, local_model_cost_map + ): model = "gemini-2.5-pro" self._token_check_patcher.stop() @@ -1294,10 +1452,16 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) - @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): + async def test_async_check_and_create_cache_skips_when_below_min_tokens( + self, mock_separate, custom_llm_provider + ): """Test that async context caching is skipped when cached content is below 1024 tokens.""" # Stop the default mock so the real token count check runs self._token_check_patcher.stop() @@ -1340,9 +1504,13 @@ class TestContextCachingEndpoints: self._token_check_patcher.start() @pytest.mark.parametrize("is_async", [False, True]) - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai"] + ) @pytest.mark.asyncio - async def test_check_and_create_cache_considers_tools_for_min_tokens(self, custom_llm_provider, is_async): + async def test_check_and_create_cache_considers_tools_for_min_tokens( + self, custom_llm_provider, is_async + ): """Test that context caching accounts for tools when validating minimum token count. Fixes #42804: When messages alone are below the threshold, but tools push the total @@ -1367,15 +1535,11 @@ class TestContextCachingEndpoints: "type": "function", "function": { "name": f"synthetic_tool_{i}", - "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " - * 8, + "description": "A very descriptive explanation of a synthetic tool designed to add tokens to the prompt cache prefix " * 8, "parameters": { "type": "object", "properties": { - f"arg_{j}": { - "type": "string", - "description": "Argument description for caching verification " * 4, - } + f"arg_{j}": {"type": "string", "description": "Argument description for caching verification " * 4} for j in range(10) }, "required": [f"arg_{j}" for j in range(5)], @@ -1404,23 +1568,19 @@ class TestContextCachingEndpoints: self.mock_client.post.return_value = mock_response self.mock_async_client.post = AsyncMock(return_value=mock_response) - with ( - patch.object( - self.context_caching, - "_get_token_and_url_context_caching", - return_value=("fake_token", "https://fake.url/cachedContents"), - ), - patch.object( - self.context_caching, - "check_cache", - return_value=None, - ), - patch.object( - self.context_caching, - "async_check_cache", - new_callable=AsyncMock, - return_value=None, - ), + with patch.object( + self.context_caching, + "_get_token_and_url_context_caching", + return_value=("fake_token", "https://fake.url/cachedContents"), + ), patch.object( + self.context_caching, + "check_cache", + return_value=None, + ), patch.object( + self.context_caching, + "async_check_cache", + new_callable=AsyncMock, + return_value=None, ): if is_async: result = await self.context_caching.async_check_and_create_cache( @@ -1464,10 +1624,13 @@ class TestContextCachingEndpoints: post_mock.assert_called_once() call_kwargs = post_mock.call_args.kwargs assert call_kwargs["json"]["tools"] == large_tools - assert call_kwargs["json"]["contents"] == [{"role": "user", "parts": [{"text": "Short system instruction."}]}] + assert call_kwargs["json"]["contents"] == [ + {"role": "user", "parts": [{"text": "Short system instruction."}]} + ] self._token_check_patcher.start() + def _model_turn_final_messages(self, final_cached_role): tool_call = { "id": "call_abc123", @@ -1514,7 +1677,9 @@ class TestContextCachingEndpoints: ] @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) - def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): + def test_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( + self, final_cached_role + ): """The cachedContents API rejects contents ending on an assistant or tool turn with HTTP 400 "Requests ending with a model turn are not supported", so the request must proceed uncached instead of failing. @@ -1547,7 +1712,9 @@ class TestContextCachingEndpoints: @pytest.mark.parametrize("final_cached_role", ["assistant", "tool", "system"]) @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn(self, final_cached_role): + async def test_async_check_and_create_cache_skips_when_cached_block_ends_on_model_turn( + self, final_cached_role + ): """Async variant: an unsupported terminal turn skips caching instead of failing.""" all_messages = self._model_turn_final_messages(final_cached_role) optional_params = self.sample_optional_params.copy() @@ -1605,8 +1772,14 @@ def test_cached_messages_end_on_supported_turn(): ) is True ) - assert cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) is False - assert cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) is False + assert ( + cached_messages_end_on_supported_turn([{"role": "tool", "tool_call_id": "x", "content": "y"}]) + is False + ) + assert ( + cached_messages_end_on_supported_turn([{"role": "function", "name": "f", "content": "y"}]) + is False + ) assert cached_messages_end_on_supported_turn([]) is False @@ -1620,9 +1793,13 @@ class TestCheckCachePagination: self.mock_client = MagicMock(spec=HTTPHandler) self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_finds_cache_on_second_page(self, mock_get_token_url, custom_llm_provider): + def test_check_cache_pagination_finds_cache_on_second_page( + self, mock_get_token_url, custom_llm_provider + ): """Test that check_cache correctly handles pagination and finds cache on second page""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1671,9 +1848,13 @@ class TestCheckCachePagination: second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"] assert "pageToken=token_page_2" in second_call_url - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): + def test_check_cache_pagination_stops_when_no_next_token( + self, mock_get_token_url, custom_llm_provider + ): """Test that check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1708,9 +1889,13 @@ class TestCheckCachePagination: assert result is None assert self.mock_client.get.call_count == 1 - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_multiple_pages(self, mock_get_token_url, custom_llm_provider): + def test_check_cache_pagination_multiple_pages( + self, mock_get_token_url, custom_llm_provider + ): """Test that check_cache correctly iterates through multiple pages""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1755,7 +1940,9 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 3 @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") async def test_async_check_cache_pagination_finds_cache_on_second_page( self, mock_get_token_url, custom_llm_provider @@ -1785,7 +1972,9 @@ class TestCheckCachePagination: } # Setup mock async client to return different responses - self.mock_async_client.get = AsyncMock(side_effect=[first_page_response, second_page_response]) + self.mock_async_client.get = AsyncMock( + side_effect=[first_page_response, second_page_response] + ) # Execute result = await self.context_caching.async_check_cache( @@ -1809,9 +1998,13 @@ class TestCheckCachePagination: assert "pageToken=token_page_2" in second_call_url @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_stops_when_no_next_token(self, mock_get_token_url, custom_llm_provider): + async def test_async_check_cache_pagination_stops_when_no_next_token( + self, mock_get_token_url, custom_llm_provider + ): """Test that async_check_cache stops pagination when no nextPageToken is present""" # Setup mock_get_token_url.return_value = ("token", "https://test-url.com") @@ -1846,21 +2039,29 @@ class TestCheckCachePagination: assert result is None assert self.mock_async_client.get.call_count == 1 - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - def test_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): + def test_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_client.get.side_effect = [create_page_response(i) for i in range(100)] + self.mock_client.get.side_effect = [ + create_page_response(i) for i in range(100) + ] result = self.context_caching.check_cache( cache_key=cache_key_to_find, @@ -1879,21 +2080,29 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 100 @pytest.mark.asyncio - @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - async def test_async_check_cache_pagination_max_pages_limit(self, mock_get_token_url, custom_llm_provider): + async def test_async_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): mock_get_token_url.return_value = ("token", "https://test-url.com") cache_key_to_find = "nonexistent_cache_key" def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], "nextPageToken": f"token_page_{page_num + 1}", } return response - self.mock_async_client.get = AsyncMock(side_effect=[create_page_response(i) for i in range(100)]) + self.mock_async_client.get = AsyncMock( + side_effect=[create_page_response(i) for i in range(100)] + ) result = await self.context_caching.async_check_cache( cache_key=cache_key_to_find, @@ -1912,6 +2121,8 @@ class TestCheckCachePagination: assert self.mock_async_client.get.call_count == 100 + + class TestVertexAIGlobalLocation: """Test global location handling in context caching.""" @@ -1937,7 +2148,9 @@ class TestVertexAIGlobalLocation: # Assert correct URL format for global expected_url = "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/cachedContents" assert url == expected_url, f"Expected {expected_url}, got {url}" - assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" + assert ( + "global-aiplatform" not in url + ), "URL should not contain 'global-aiplatform' prefix" def test_regional_location_url_construction_v1(self): """Test that regional location uses correct URL (with location prefix) for v1 API.""" @@ -1980,11 +2193,11 @@ class TestVertexAIGlobalLocation: ) # Assert correct URL format for global with beta API - expected_url = ( - "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" - ) + expected_url = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" assert url == expected_url, f"Expected {expected_url}, got {url}" - assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix" + assert ( + "global-aiplatform" not in url + ), "URL should not contain 'global-aiplatform' prefix" def test_gemini_context_caching_with_custom_api_base_passes_model(self): """Gemini context caching with custom api_base must pass model to _check_custom_proxy. @@ -2079,6 +2292,8 @@ class TestContextCachingMultiRegionUrls: _STORAGE_TEST_MODEL: Final = "gemini-2.5-pro" + + _STORAGE_TEST_CACHE_NAME: Final = "cachedContents/storage-cost" diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 726fb4c33c3..30a04290780 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -189,9 +189,7 @@ def test_response_cost_calculator_keeps_optional_params_out_of_hidden_params(): assert optional_params["aws_session_token"] == "session-secret" -def test_embedding_success_logging_and_spend_log_carry_no_forwarded_credentials( - monkeypatch: pytest.MonkeyPatch, -) -> None: +def test_embedding_success_logging_and_spend_log_carry_no_forwarded_credentials(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server from litellm.proxy.spend_tracking.spend_tracking_utils import _get_proxy_server_request_for_spend_logs_payload @@ -240,6 +238,10 @@ def test_embedding_success_logging_and_spend_log_carry_no_forwarded_credentials( assert logging_obj.optional_params["extra_headers"] == {"x-goog-api-key": "goog-secret"} + + + + def test_realtime_stream_combines_text_and_audio_token_details(): """Realtime response.done usage with input_token_details / output_token_details.""" from litellm.cost_calculator import RealtimeAPITokenUsageProcessor @@ -1354,6 +1356,8 @@ def test_bedrock_cost_calculator_comparison_with_without_cache(): print(f"Cost with cache: {cost_with_cache}") + + def test_gemini_25_explicit_caching_cost_direct_usage(): """ Test that Gemini 2.5 models correctly calculate costs with explicit caching. @@ -1990,6 +1994,8 @@ def test_cost_margin_with_discount(monkeypatch): print(f" - Expected: ${expected_cost:.6f}") + + def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_map): """Test that completion_cost extracts service_tier from completion_response object.""" from litellm import completion_cost @@ -2739,6 +2745,8 @@ def test_gemini_without_cache_tokens_details(): print("✅ Gemini without cacheTokensDetails works correctly") + + def test_additional_costs_only_for_azure_ai(_local_model_cost_map): """ Test that _get_additional_costs is only called for azure_ai provider. @@ -3283,7 +3291,9 @@ def test_cost_per_token_resolves_per_second_rate_precedence( model: Final = "test-chat-per-second-rate-precedence" entry: Final = {**pricing_fields, "litellm_provider": "together_ai", "mode": "chat"} - litellm.register_model(model_cost={model: entry}) + litellm.register_model( + model_cost={model: entry} + ) assert cost_per_token( model=model, @@ -4051,10 +4061,7 @@ def test_completion_cost_region_without_its_own_row_prices_mantle_claude_from_th custom_llm_provider="bedrock_mantle", region_name="us-east-1", ) == pytest.approx(expected), deployment - assert ( - litellm.get_model_info(f"bedrock_mantle/us-east-1/{model}", "bedrock_mantle")["key"] - == f"bedrock_mantle/{model}" - ) + assert litellm.get_model_info(f"bedrock_mantle/us-east-1/{model}", "bedrock_mantle")["key"] == f"bedrock_mantle/{model}" @pytest.mark.parametrize("model", ["anthropic.claude-opus-5-5", "anthropic.claude-sonnet-5-5"]) @@ -4978,7 +4985,9 @@ def test_xai_batch_tier_discounts_the_long_context_rate_like_the_flat_batch_rate assert info[f"{prefix}_above_200k_tokens_batches"] < info[f"{prefix}_above_200k_tokens"] -@pytest.mark.parametrize(("prompt_tokens", "tier"), [(200_000, "_above_200k_tokens_batches"), (199_999, "_batches")]) +@pytest.mark.parametrize( + ("prompt_tokens", "tier"), [(200_000, "_above_200k_tokens_batches"), (199_999, "_batches")] +) def test_xai_batch_cost_calculator_bills_the_200k_batch_tier_inclusively( _local_model_cost_map: None, prompt_tokens: int, tier: str ) -> None: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index cb99e178dc4..9201eaa4e2a 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -63,7 +63,7 @@ from litellm.types.utils import ( bedrock_batch_litellm_params, ) from litellm.types.videos.main import VideoObject -from litellm.utils import ( +from litellm.utils import( _invalidate_model_cost_lowercase_map, CustomStreamWrapper, filter_out_litellm_params, @@ -3165,7 +3165,9 @@ class TestExtraBodyCannotOverrideModel: } ] - untouched = litellm.get_optional_params(model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools()) + untouched = litellm.get_optional_params( + model="my-vllm-model", custom_llm_provider="hosted_vllm", tools=tools() + ) assert untouched["tools"][0]["function"]["custom_marker"] == "LEAK", untouched result = litellm.get_optional_params( @@ -4156,9 +4158,7 @@ def test_is_prompt_caching_valid_prompt_stops_counting_once_the_minimum_is_reach assert is_prompt_caching_valid_prompt(model="claude-opus-4-8", messages=long_prompt, min_token_count=1024) is True assert sum(counted_messages) < len(long_prompt), sum(counted_messages) - full_count = litellm.token_counter( - model="claude-opus-4-8", messages=long_prompt, use_default_image_token_count=True - ) + full_count = litellm.token_counter(model="claude-opus-4-8", messages=long_prompt, use_default_image_token_count=True) assert ( is_prompt_caching_valid_prompt(model="claude-opus-4-8", messages=long_prompt, min_token_count=full_count) is True @@ -4621,13 +4621,7 @@ _SUCCESS_RESPONSES_BY_CALL_TYPE: Final = ( pytest.param(EmbeddingResponse(model="text-embedding-3-small"), CallTypes.aembedding, id="embedding"), pytest.param( ResponsesAPIResponse( - id="resp_abc", - created_at=1, - output=[], - parallel_tool_calls=False, - tool_choice="auto", - tools=[], - model="gpt-5.6", + id="resp_abc", created_at=1, output=[], parallel_tool_calls=False, tool_choice="auto", tools=[], model="gpt-5.6" ), CallTypes.aresponses, id="responses", @@ -4655,9 +4649,7 @@ async def test_success_deployment_hook_raising_keeps_response_and_runs_later_hoo assert result is response assert second_hook.seen_responses == (response,) - failure_logs: Final = tuple( - r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message - ) + failure_logs: Final = tuple(r for r in caplog.records if "async_post_call_success_deployment_hook error" in r.message) assert len(failure_logs) == 1 assert "_ChatShapedSuccessDeploymentHook" in failure_logs[0].message assert str(call_type) in failure_logs[0].message @@ -4794,13 +4786,17 @@ async def test_wrapper_async_logs_converted_responses_stream_with_standard_loggi ({"metadata": None}, True, False), ], ) -def test_is_litellm_router_call_is_async_aware(kwargs: Mapping[str, object], is_async: bool, expected: bool) -> None: +def test_is_litellm_router_call_is_async_aware( + kwargs: Mapping[str, object], is_async: bool, expected: bool +) -> None: assert _is_litellm_router_call(kwargs, is_async=is_async) is expected @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) -async def test_router_aresponses_does_not_run_sdk_retries(monkeypatch: pytest.MonkeyPatch, stream: bool) -> None: +async def test_router_aresponses_does_not_run_sdk_retries( + monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() model_list: Final = [ @@ -4814,7 +4810,9 @@ async def test_router_aresponses_does_not_run_sdk_retries(monkeypatch: pytest.Mo }, } ] - router: Final = litellm.Router(model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True) + router: Final = litellm.Router( + model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True + ) try: with respx.mock(assert_all_called=True) as respx_mock: @@ -4850,7 +4848,9 @@ def test_router_responses_keeps_sdk_retries_for_sync_router_call( }, } ] - router: Final = litellm.Router(model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True) + router: Final = litellm.Router( + model_list=model_list, num_retries=0, retry_after=0, disable_cooldowns=True + ) try: with respx.mock(assert_all_called=True) as respx_mock: @@ -6694,9 +6694,7 @@ def test_function_setup_logs_the_search_query_edit_prompt_and_ocr_document_summa @pytest.mark.parametrize("original_function", ("atext_completion", "text_completion")) -def test_function_setup_without_a_prompt_leaves_the_missing_prompt_to_request_validation( - original_function: str, -) -> None: +def test_function_setup_without_a_prompt_leaves_the_missing_prompt_to_request_validation(original_function: str) -> None: assert _logged_request_messages(original_function, model="gpt-4o") is None @@ -6738,7 +6736,6 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) - @pytest.fixture(scope="function") def setup_and_teardown(): """ @@ -6752,17 +6749,14 @@ def setup_and_teardown(): loop.close() asyncio.set_event_loop(None) - MODEL: Final = "anthropic/claude-haiku-4-5" - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_none(): """Test that None is returned as-is.""" result = validate_chat_completion_tool_choice(None, model=MODEL) assert result is None - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_string(): """Test that string values are returned as-is.""" @@ -6770,7 +6764,6 @@ def test_validate_tool_choice_string(): assert validate_chat_completion_tool_choice("none", model=MODEL) == "none" assert validate_chat_completion_tool_choice("required", model=MODEL) == "required" - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_standard_dict(): """Test standard OpenAI format with function.""" @@ -6778,7 +6771,6 @@ def test_validate_tool_choice_standard_dict(): result = validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert result == tool_choice - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_cursor_format(): """Cursor IDE format {"type": "auto"} is unwrapped to the bare string.""" @@ -6786,7 +6778,6 @@ def test_validate_tool_choice_cursor_format(): assert validate_chat_completion_tool_choice({"type": "none"}, model=MODEL) == "none" assert validate_chat_completion_tool_choice({"type": "required"}, model=MODEL) == "required" - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize( "tool_choice", @@ -6808,7 +6799,6 @@ def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): assert exc_info.value.status_code == 400 assert exc_info.value.model == MODEL - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("tool_choice", [123, []]) def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): @@ -6819,7 +6809,6 @@ def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert exc_info.value.status_code == 400 - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_validate_tool_choice_without_model_is_still_a_400(): """Callers that predate the model argument keep getting a 400, with an empty model on the error.""" @@ -6828,14 +6817,12 @@ def test_validate_tool_choice_without_model_is_still_a_400(): assert exc_info.value.status_code == 400 assert exc_info.value.model == "" - @pytest.fixture() def _vcr_outcome_gate_local_testing(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) - @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -6888,7 +6875,6 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() - _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -6909,7 +6895,6 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } - @pytest.fixture(scope="module") def setup_and_teardown_local_testing(): """ @@ -6932,7 +6917,6 @@ def setup_and_teardown_local_testing(): litellm.in_memory_llm_clients_cache.flush_cache() yield - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -6958,7 +6942,6 @@ def test_vertex_projects(): assert "vertex_ai_project" in optional_params assert "vertex_ai_location" in optional_params - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -6974,7 +6957,6 @@ def test_bedrock_embed_v2_regular(): print(f"received optional_params: {optional_params}") assert optional_params == {"dimensions": 512} - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -6993,7 +6975,6 @@ def test_bedrock_embed_v2_with_drop_params(): print(f"received optional_params: {optional_params}") assert optional_params == {"dimensions": 512, "embeddingTypes": ["binary"]} - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7015,7 +6996,6 @@ def test_openai_non_text_embedding_3_with_allowed_openai_params(): print(f"received optional_params: {optional_params}") assert optional_params.get("dimensions") == 1024 - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7042,7 +7022,6 @@ def test_openai_non_text_embedding_3_without_allowed_openai_params_raises(): finally: litellm.drop_params = prev_drop_params - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7071,7 +7050,6 @@ def test_openai_non_text_embedding_3_drop_params_per_call(): finally: litellm.drop_params = prev_drop_params - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -7099,14 +7077,12 @@ def test_openai_non_text_embedding_3_drop_params_global(): finally: litellm.drop_params = prev_drop_params - @pytest.fixture() def _vcr_outcome_gate_search_tests(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) - @pytest.mark.usefixtures("_vcr_outcome_gate_search_tests") def test_search_tool_name_in_all_litellm_params(): """ @@ -7116,7 +7092,6 @@ def test_search_tool_name_in_all_litellm_params(): """ assert "search_tool_name" in all_litellm_params - @pytest.mark.usefixtures("_vcr_outcome_gate_search_tests") def test_filter_out_search_tool_name(): """ @@ -7143,7 +7118,6 @@ def test_filter_out_search_tool_name(): assert filtered["query"] == "latest ai developments" assert filtered["max_results"] == 5 - @pytest.mark.asyncio async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None: """Chat over the Responses bridge exits two @client wrappers with one logging object. Issue From c231b7f530f122b9ff6efd377fdd2d1b30904932 Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:49:31 +0900 Subject: [PATCH 03/10] fix(vertex_ai): keep token pricing for storage-only overrides and price storage from the underlying model A deployment that sets only cache_storage_cost_per_token_per_hour no longer turns on custom pricing, so its input and output tokens keep the model's rates. The storage rate is read from the deployment entry first, then from the cost map entry of the underlying model rather than the deployment id, so custom token prices without a storage rate still bill storage. --- litellm/cost_calculator.py | 27 ++++++-- litellm/litellm_core_utils/litellm_logging.py | 5 +- litellm/types/utils.py | 4 +- .../context_caching/test_storage_cost.py | 12 ++-- .../test_vertex_ai_context_caching.py | 36 +++++++--- tests/unit/test_cost_calculator.py | 66 ++++++++++++++++++- tests/unit/test_utils.py | 2 +- 7 files changed, 128 insertions(+), 24 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e1eeb10ab3d..7ed8a20922b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -301,20 +301,30 @@ def _get_additional_costs( def _context_cache_storage_cost( litellm_logging_obj: LitellmLoggingObject | None, - model: str, + *, + requested_model: str | None, + completion_response: object | None, + base_model: str | None, custom_llm_provider: str | None, - custom_pricing: bool | None, router_model_id: str | None, + region_name: str | None, ) -> float: token_hours: Final = litellm_logging_obj.context_cache_storage_token_hours if litellm_logging_obj else 0.0 if not isinstance(token_hours, int | float) or token_hours <= 0: return 0.0 + underlying_model: Final = select_model_name_for_cost_calc( + model=requested_model, + completion_response=completion_response, + base_model=base_model, + custom_llm_provider=custom_llm_provider, + region_name=region_name, + ) rate: Final = next( ( price for source in ( - _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id), - _cost_map_model_info(model, custom_llm_provider), + _deployment_model_info(litellm_logging_obj, custom_pricing=True, router_model_id=router_model_id), + _cost_map_model_info(underlying_model, custom_llm_provider) if underlying_model else None, ) if source is not None and isinstance(price := source.get("cache_storage_cost_per_token_per_hour"), int | float) @@ -1480,6 +1490,7 @@ def completion_cost( region_name=region_name, ) + requested_model: Final = model potential_model_names: Final = [ selected_model, _get_response_model(completion_response), @@ -1873,7 +1884,13 @@ def completion_cost( else: additional_costs = None cache_storage_cost: Final = _context_cache_storage_cost( - litellm_logging_obj, model, custom_llm_provider, custom_pricing, router_model_id + litellm_logging_obj, + requested_model=requested_model, + completion_response=completion_response, + base_model=base_model, + custom_llm_provider=custom_llm_provider, + router_model_id=router_model_id, + region_name=region_name, ) if cache_storage_cost > 0: additional_costs = {**(additional_costs or {}), "cache_storage_cost": cache_storage_cost} diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 766cceceff7..82c382d451e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -144,6 +144,7 @@ from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.rerank import RerankResponse from litellm.types.utils import ( DEPLOYMENT_SCOPED_PRICING_FIELDS, + SEPARATE_CHARGE_PRICING_FIELDS, CachingDetails, CallTypes, CostBreakdown, @@ -323,7 +324,9 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: ### GLOBAL VARIABLES ### # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys -_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) +_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = ( + frozenset(CustomPricingLiteLLMParams.model_fields.keys()) - SEPARATE_CHARGE_PRICING_FIELDS +) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 791ad9fddcd..82460ba9414 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -319,7 +319,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_batches: ReadOnly[float | None] cache_creation_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_creation_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] - cache_storage_cost_per_token_per_hour: ReadOnly[float | None] # explicit context cache storage, e.g. Vertex AI + cache_storage_cost_per_token_per_hour: ReadOnly[float | None] # Smallest prefix this model will actually cache, whatever caching mechanism its provider uses. # Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT. prompt_cache_min_tokens: int | None @@ -3927,6 +3927,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) +SEPARATE_CHARGE_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"cache_storage_cost_per_token_per_hour"}) + DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"}) SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( diff --git a/tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py b/tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py index b3bfb02b4e1..5fd3bbad9cc 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_storage_cost.py @@ -1,3 +1,5 @@ +from typing import Final + import pytest from litellm.llms.vertex_ai.context_caching.storage_cost import ( @@ -30,18 +32,20 @@ from litellm.llms.vertex_ai.context_caching.storage_cost import ( ("not-an-object", "3600s", 0.0), ], ) -def test_token_hours_follow_reported_lifetime_then_requested_ttl(created, requested_ttl, expected): +def test_token_hours_follow_reported_lifetime_then_requested_ttl( + created: object, requested_ttl: str | None, expected: float +) -> None: assert context_cache_storage_token_hours(created, requested_ttl) == pytest.approx(expected) -def test_token_hours_default_to_one_hour_without_lifetime_or_ttl(): +def test_token_hours_default_to_one_hour_without_lifetime_or_ttl() -> None: # Source: https://cloud.google.com/vertex-ai/generative-ai/docs/context-cache/context-cache-create # "The default expiration time of a context cache is 60 minutes after it's created" (read 2026-10-07) assert context_cache_storage_token_hours({"usageMetadata": {"totalTokenCount": 5000}}, None) == 5000.0 @pytest.mark.parametrize("requested_ttl", ["", "s", "abc", "10m", "-5s"]) -def test_unparseable_requested_ttl_falls_back_to_default_lifetime(requested_ttl): - created = {"usageMetadata": {"totalTokenCount": 3600}} +def test_unparseable_requested_ttl_falls_back_to_default_lifetime(requested_ttl: str) -> None: + created: Final = {"usageMetadata": {"totalTokenCount": 3600}} assert context_cache_storage_token_hours(created, requested_ttl) == context_cache_storage_token_hours(created, None) 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 77907e0ca19..bf9f73aa08e 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,8 +1,10 @@ -from typing import Final +from collections.abc import Mapping +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from typing_extensions import ReadOnly, TypedDict import litellm @@ -12,6 +14,7 @@ from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( ContextCachingEndpoints, ) +from litellm.types.llms.openai import AllMessageValues @pytest.fixture @@ -2313,7 +2316,7 @@ def _storage_logging_obj() -> Logging: return logging_obj -def _storage_messages() -> list: +def _storage_messages() -> list[AllMessageValues]: return [ { "role": "system", @@ -2323,7 +2326,7 @@ def _storage_messages() -> list: ] -def _json_response(method: str, payload: dict) -> httpx.Response: +def _json_response(method: str, payload: Mapping[str, object]) -> httpx.Response: return httpx.Response(200, json=payload, request=httpx.Request(method, "https://example.invalid")) @@ -2336,7 +2339,22 @@ _CREATED_CACHE: Final = { } -def _create_kwargs(logging_obj: Logging, client: object) -> dict: +class _CreateCacheKwargs(TypedDict): + messages: ReadOnly[list[AllMessageValues]] + optional_params: ReadOnly[dict[str, object]] + api_key: ReadOnly[str] + api_base: ReadOnly[str | None] + model: ReadOnly[str] + client: ReadOnly[MagicMock] + timeout: ReadOnly[float] + logging_obj: ReadOnly[Logging] + custom_llm_provider: ReadOnly[Literal["gemini"]] + vertex_project: ReadOnly[str | None] + vertex_location: ReadOnly[str | None] + vertex_auth_header: ReadOnly[str | None] + + +def _create_kwargs(logging_obj: Logging, client: MagicMock) -> _CreateCacheKwargs: return { "messages": _storage_messages(), "optional_params": {}, @@ -2353,7 +2371,7 @@ def _create_kwargs(logging_obj: Logging, client: object) -> dict: } -def test_creating_a_cache_records_its_storage_token_hours(local_model_cost_map): +def test_creating_a_cache_records_its_storage_token_hours(local_model_cost_map: None) -> None: logging_obj: Final = _storage_logging_obj() client: Final = MagicMock(spec=HTTPHandler) client.get.return_value = _json_response("GET", {}) @@ -2366,7 +2384,7 @@ def test_creating_a_cache_records_its_storage_token_hours(local_model_cost_map): @pytest.mark.asyncio -async def test_async_creating_a_cache_records_its_storage_token_hours(local_model_cost_map): +async def test_async_creating_a_cache_records_its_storage_token_hours(local_model_cost_map: None) -> None: logging_obj: Final = _storage_logging_obj() client: Final = MagicMock(spec=AsyncHTTPHandler) client.get = AsyncMock(return_value=_json_response("GET", {})) @@ -2380,11 +2398,11 @@ async def test_async_creating_a_cache_records_its_storage_token_hours(local_mode assert logging_obj.context_cache_storage_token_hours == pytest.approx(7200 * 0.5) -def _listing_with(created_body: dict) -> dict: +def _listing_with(created_body: Mapping[str, object]) -> Mapping[str, object]: return {"cachedContents": [{"name": _STORAGE_TEST_CACHE_NAME, "displayName": created_body["displayName"]}]} -def test_reusing_an_existing_cache_records_no_storage(local_model_cost_map): +def test_reusing_an_existing_cache_records_no_storage(local_model_cost_map: None) -> None: client: Final = MagicMock(spec=HTTPHandler) client.get.return_value = _json_response("GET", {}) client.post.return_value = _json_response("POST", _CREATED_CACHE) @@ -2400,7 +2418,7 @@ def test_reusing_an_existing_cache_records_no_storage(local_model_cost_map): @pytest.mark.asyncio -async def test_async_reusing_an_existing_cache_records_no_storage(local_model_cost_map): +async def test_async_reusing_an_existing_cache_records_no_storage(local_model_cost_map: None) -> None: client: Final = MagicMock(spec=AsyncHTTPHandler) client.get = AsyncMock(return_value=_json_response("GET", {})) client.post = AsyncMock(return_value=_json_response("POST", _CREATED_CACHE)) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 30a04290780..6abaee97b92 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -1,6 +1,6 @@ import datetime import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Final, cast @@ -5844,8 +5844,11 @@ def test_completion_cost_bills_base_when_gemini_serves_on_demand( _STORAGE_RATE_KEY: Final = "cache_storage_cost_per_token_per_hour" +_FIXED_START_TIME: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + + @pytest.fixture -def _storage_cost_map(_local_model_cost_map: None): +def _storage_cost_map(_local_model_cost_map: None) -> Iterator[None]: litellm.get_model_info.cache_clear() yield litellm.get_model_info.cache_clear() @@ -5874,7 +5877,7 @@ def _cache_storage_logging_obj(model: str, token_hours: float) -> Logging: messages=[{"role": "user", "content": "Hello"}], stream=False, call_type="completion", - start_time=datetime.datetime.now(), + start_time=_FIXED_START_TIME, litellm_call_id="context-cache-storage", function_id="f", ) @@ -5937,3 +5940,60 @@ def test_context_cache_storage_uses_the_deployment_storage_rate(_storage_cost_ma without_storage: Final = _chat_cost(model, _cache_storage_logging_obj(model, token_hours=0.0), **pricing) assert with_storage - without_storage == pytest.approx(500.0 * deployment_rate) + + +_STORAGE_DEPLOYMENT_MODEL: Final = "gemini-2.5-flash" + + +def _routed_request_cost(deployment_pricing: Mapping[str, float], token_hours: float) -> float | None: + from litellm import Router + + router: Final = Router( + model_list=[ + { + "model_name": "storage-cost-deployment", + "litellm_params": {"model": f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}", **deployment_pricing}, + } + ] + ) + deployment_id: Final = router.model_list[0]["model_info"]["id"] + return _vertex_request_cost({**deployment_pricing, "metadata": {"model_info": {"id": deployment_id}}}, token_hours) + + +def _vertex_request_cost(litellm_params: Mapping[str, object], token_hours: float) -> float | None: + logging_obj: Final = _cache_storage_logging_obj(_STORAGE_DEPLOYMENT_MODEL, token_hours) + logging_obj.update_environment_variables( + model=_STORAGE_DEPLOYMENT_MODEL, + litellm_params=dict(litellm_params), + optional_params={}, + custom_llm_provider="vertex_ai", + ) + usage: Final = Usage(prompt_tokens=1_000, completion_tokens=100, total_tokens=1_100) + return logging_obj.response_cost_calculator(ModelResponse(model=_STORAGE_DEPLOYMENT_MODEL, usage=usage)) + + +def test_custom_token_prices_without_a_storage_rate_bill_storage_at_the_underlying_model_rate( + _storage_cost_map: None, +) -> None: + cost_map_rate: Final = litellm.get_model_info(f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}")[_STORAGE_RATE_KEY] + token_prices: Final = {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6} + + with_storage: Final = _routed_request_cost(token_prices, token_hours=500.0) + without_storage: Final = _routed_request_cost(token_prices, token_hours=0.0) + + assert cost_map_rate is not None and cost_map_rate > 0 + assert without_storage == pytest.approx(1_000 * 1e-6 + 100 * 2e-6) + assert with_storage == pytest.approx(without_storage + 500.0 * cost_map_rate) + + +def test_a_storage_only_override_keeps_the_underlying_token_prices(_storage_cost_map: None) -> None: + deployment_rate: Final = litellm.get_model_info(f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}")[_STORAGE_RATE_KEY] * 3 + storage_only: Final = {_STORAGE_RATE_KEY: deployment_rate} + + no_override: Final = _vertex_request_cost({}, token_hours=0.0) + without_storage: Final = _routed_request_cost(storage_only, token_hours=0.0) + with_storage: Final = _routed_request_cost(storage_only, token_hours=500.0) + + assert no_override is not None and no_override > 0 + assert without_storage == pytest.approx(no_override) + assert with_storage == pytest.approx(no_override + 500.0 * deployment_rate) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9201eaa4e2a..41400679068 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -7182,7 +7182,7 @@ async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: assert counting_logger.logged_results[0] is inner_result -def test_get_model_info_surfaces_context_cache_storage_rate(monkeypatch): +def test_get_model_info_surfaces_context_cache_storage_rate(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) litellm.get_model_info.cache_clear() From 2ce6744ea5fad56fdc97bed32d7db8777ed62023 Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:12:06 +0900 Subject: [PATCH 04/10] fix(vertex_ai): keep basedpyright counts flat for cache storage pricing Narrow the arguments passed to the storage cost helper, add the storage charge to the total outside the additional_costs sum, and type the new test helpers so the changed files add no reportUnknownArgumentType or reportGeneralTypeIssues errors. Raise the reportUnknownArgumentType limit by 97 for the new cache_storage_cost_per_token_per_hour pricing field: like every pricing field, it adds one diagnostic at each existing untyped GenericLiteLLMParams(**kwargs) call site (same approach as #43764). --- basedpyright-code-budget.json | 2 +- litellm/cost_calculator.py | 18 ++++++++++++------ tests/unit/test_cost_calculator.py | 23 +++++++++++++++++------ tests/unit/test_utils.py | 3 ++- 4 files changed, 32 insertions(+), 14 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 92dc89eb0b8..0b1a357fffd 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44802 + "limit": 44899 }, "reportUnknownLambdaType": { "limit": 109 diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 2e75be27330..dfe63dc28f6 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1883,17 +1883,20 @@ def completion_cost( ) else: additional_costs = None - cache_storage_cost: Final = _context_cache_storage_cost( + _storage_response: ModelResponse | None = ( + completion_response if isinstance(completion_response, ModelResponse) else None + ) + _storage_provider: str | None = custom_llm_provider if isinstance(custom_llm_provider, str) else None + _storage_region: str | None = region_name if isinstance(region_name, str) else None + cache_storage_cost: float = _context_cache_storage_cost( litellm_logging_obj, requested_model=requested_model, - completion_response=completion_response, + completion_response=_storage_response, base_model=base_model, - custom_llm_provider=custom_llm_provider, + custom_llm_provider=_storage_provider, router_model_id=router_model_id, - region_name=region_name, + region_name=_storage_region, ) - if cache_storage_cost > 0: - additional_costs = {**(additional_costs or {}), "cache_storage_cost": cache_storage_cost} _final_cost = prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar cost_for_built_in_tools = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( @@ -1906,6 +1909,9 @@ def completion_cost( _final_cost += cost_for_built_in_tools if additional_costs: _final_cost += sum(additional_costs.values()) + if cache_storage_cost > 0: + _final_cost += cache_storage_cost + additional_costs = {**(additional_costs or {}), "cache_storage_cost": cache_storage_cost} original_cost = _final_cost if litellm.cost_discount_config: diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 3cccbcbbf4d..02f1cee9c12 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -5832,19 +5832,27 @@ def _storage_cost_map(_local_model_cost_map: None) -> Iterator[None]: def _first_chat_model(provider_prefix: str, *, with_storage_rate: bool) -> str: + names: Final = [name for name in litellm.model_cost if isinstance(name, str)] return next( name - for name, entry in litellm.model_cost.items() - if _is_priced_chat_entry(entry, provider_prefix) and (_STORAGE_RATE_KEY in entry) == with_storage_rate + for name in names + if _is_priced_chat_entry(_cost_map_entry(name), provider_prefix, with_storage_rate=with_storage_rate) ) -def _is_priced_chat_entry(entry: object, provider_prefix: str) -> bool: +def _cost_map_entry(name: str) -> object: + entry: object = litellm.model_cost.get(name) + return entry + + +def _is_priced_chat_entry(entry: object, provider_prefix: str, *, with_storage_rate: bool) -> bool: return ( isinstance(entry, dict) and entry.get("mode") == "chat" - and str(entry.get("litellm_provider", "")).startswith(provider_prefix) + and isinstance(provider := entry.get("litellm_provider"), str) + and provider.startswith(provider_prefix) and entry.get("input_cost_per_token") is not None + and (_STORAGE_RATE_KEY in entry) == with_storage_rate ) @@ -5934,6 +5942,7 @@ def _routed_request_cost(deployment_pricing: Mapping[str, float], token_hours: f ] ) deployment_id: Final = router.model_list[0]["model_info"]["id"] + assert isinstance(deployment_id, str) return _vertex_request_cost({**deployment_pricing, "metadata": {"model_info": {"id": deployment_id}}}, token_hours) @@ -5964,8 +5973,10 @@ def test_custom_token_prices_without_a_storage_rate_bill_storage_at_the_underlyi def test_a_storage_only_override_keeps_the_underlying_token_prices(_storage_cost_map: None) -> None: - deployment_rate: Final = litellm.get_model_info(f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}")[_STORAGE_RATE_KEY] * 3 - storage_only: Final = {_STORAGE_RATE_KEY: deployment_rate} + cost_map_rate: Final = litellm.get_model_info(f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}")[_STORAGE_RATE_KEY] + assert isinstance(cost_map_rate, float) + deployment_rate: Final = cost_map_rate * 3 + storage_only: Final[dict[str, float]] = {_STORAGE_RATE_KEY: deployment_rate} no_override: Final = _vertex_request_cost({}, token_hours=0.0) without_storage: Final = _routed_request_cost(storage_only, token_hours=0.0) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a9aa66f9565..1be5bf8a4ea 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -7216,7 +7216,8 @@ async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: def test_get_model_info_surfaces_context_cache_storage_rate(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + local_cost_map: dict[str, object] = litellm.get_model_cost_map(url="") + monkeypatch.setattr(litellm, "model_cost", local_cost_map) litellm.get_model_info.cache_clear() priced: Final = { name: entry["cache_storage_cost_per_token_per_hour"] From 2ef85f8a4addb990545cc8228cdd7143e947b271 Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:04:07 +0900 Subject: [PATCH 05/10] test: annotate new cost-map test locals as Final --- tests/unit/test_cost_calculator.py | 2 +- tests/unit/test_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 02f1cee9c12..68764e5b431 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -5841,7 +5841,7 @@ def _first_chat_model(provider_prefix: str, *, with_storage_rate: bool) -> str: def _cost_map_entry(name: str) -> object: - entry: object = litellm.model_cost.get(name) + entry: Final[object] = litellm.model_cost.get(name) return entry diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 1be5bf8a4ea..fc3d4a579f8 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -7216,7 +7216,7 @@ async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: def test_get_model_info_surfaces_context_cache_storage_rate(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - local_cost_map: dict[str, object] = litellm.get_model_cost_map(url="") + local_cost_map: Final[Mapping[str, object]] = litellm.get_model_cost_map(url="") monkeypatch.setattr(litellm, "model_cost", local_cost_map) litellm.get_model_info.cache_clear() priced: Final = { From 45c52ad8d4fcbc903e276602ef36bcb4c2b47ede Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:19:30 +0900 Subject: [PATCH 06/10] refactor(vertex_ai): take the deployment storage rate from model_info The deployment-level `cache_storage_cost_per_token_per_hour` override moves from `litellm_params` to `model_info`, next to `cost_per_ptu_per_hour`. It is a separate per-hour charge rather than a token price, so it no longer joins `CustomPricingLiteLLMParams`: it never switches on custom pricing, and every untyped `GenericLiteLLMParams(**kwargs)` call site stops gaining a diagnostic, which keeps basedpyright at its merge-base counts now that the budget file is gone. The field is kept off the shared `{provider}/{model}` cost-map key so one deployment's rate never reaches its siblings. --- litellm/cost_calculator.py | 11 +++++++-- litellm/litellm_core_utils/litellm_logging.py | 5 +--- litellm/types/router.py | 1 + litellm/types/utils.py | 4 ++-- tests/unit/test_cost_calculator.py | 24 ++++++++++++++----- .../unit/test_router_model_cost_isolation.py | 8 +++++-- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 ++--- 7 files changed, 39 insertions(+), 20 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a15c07bd77d..2a01135216d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -299,6 +299,10 @@ def _get_additional_costs( return None +def _recorded_cache_storage_token_hours(litellm_logging_obj: LitellmLoggingObject | None) -> object: + return litellm_logging_obj.context_cache_storage_token_hours if litellm_logging_obj else 0.0 + + def _context_cache_storage_cost( litellm_logging_obj: LitellmLoggingObject | None, *, @@ -309,7 +313,7 @@ def _context_cache_storage_cost( router_model_id: str | None, region_name: str | None, ) -> float: - token_hours: Final = litellm_logging_obj.context_cache_storage_token_hours if litellm_logging_obj else 0.0 + token_hours: Final = _recorded_cache_storage_token_hours(litellm_logging_obj) if not isinstance(token_hours, int | float) or token_hours <= 0: return 0.0 underlying_model: Final = select_model_name_for_cost_calc( @@ -1935,7 +1939,10 @@ def completion_cost( _final_cost += sum(additional_costs.values()) if cache_storage_cost > 0: _final_cost += cache_storage_cost - additional_costs = {**(additional_costs or {}), "cache_storage_cost": cache_storage_cost} + if additional_costs is None: + additional_costs = {"cache_storage_cost": cache_storage_cost} + else: + additional_costs["cache_storage_cost"] = cache_storage_cost original_cost = _final_cost if litellm.cost_discount_config: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 379d202a65f..f034107b886 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -144,7 +144,6 @@ from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.rerank import RerankResponse from litellm.types.utils import ( DEPLOYMENT_SCOPED_PRICING_FIELDS, - SEPARATE_CHARGE_PRICING_FIELDS, CachingDetails, CallTypes, CostBreakdown, @@ -324,9 +323,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: ### GLOBAL VARIABLES ### # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys -_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = ( - frozenset(CustomPricingLiteLLMParams.model_fields.keys()) - SEPARATE_CHARGE_PRICING_FIELDS -) +_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") diff --git a/litellm/types/router.py b/litellm/types/router.py index a66c4571b39..4da354fd818 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -274,6 +274,7 @@ class ModelInfo(MirroredPricingParams): cost_per_ptu_per_hour: float | None = None ptu_effective_from: datetime.datetime | None = None ptu_effective_to: datetime.datetime | None = None + cache_storage_cost_per_token_per_hour: float | None = None # when tag-based routing's "!" or "&" constraints eliminate every deployment # in this model group, fall back to the default-tagged pool instead of diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6fcde3e5ec3..352b8ad4cb3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3857,7 +3857,6 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_creation_input_token_cost_above_100k_tokens_batches: float | None = None cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None - cache_storage_cost_per_token_per_hour: float | None = None cache_read_input_audio_token_cost: float | None = None cache_read_input_image_token_cost: float | None = None input_cost_per_character_above_128k_tokens: float | None = None @@ -3947,7 +3946,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) -SEPARATE_CHARGE_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"cache_storage_cost_per_token_per_hour"}) +DEPLOYMENT_SCOPED_SEPARATE_CHARGE_FIELDS: Final[frozenset[str]] = frozenset({"cache_storage_cost_per_token_per_hour"}) DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"}) @@ -3955,6 +3954,7 @@ SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) - frozenset(CustomPricingLiteLLMParams.model_fields) - DEPLOYMENT_SCOPED_PRICING_FIELDS + - DEPLOYMENT_SCOPED_SEPARATE_CHARGE_FIELDS - DEPLOYMENT_SCOPED_CAPABILITY_FIELDS ) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 68764e5b431..02570a38d79 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -5930,20 +5930,23 @@ def test_context_cache_storage_uses_the_deployment_storage_rate(_storage_cost_ma _STORAGE_DEPLOYMENT_MODEL: Final = "gemini-2.5-flash" -def _routed_request_cost(deployment_pricing: Mapping[str, float], token_hours: float) -> float | None: +def _routed_request_cost( + token_prices: Mapping[str, float], token_hours: float, *, storage_rate: float | None = None +) -> float | None: from litellm import Router router: Final = Router( model_list=[ { "model_name": "storage-cost-deployment", - "litellm_params": {"model": f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}", **deployment_pricing}, + "litellm_params": {"model": f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}", **token_prices}, + "model_info": {} if storage_rate is None else {_STORAGE_RATE_KEY: storage_rate}, } ] ) deployment_id: Final = router.model_list[0]["model_info"]["id"] assert isinstance(deployment_id, str) - return _vertex_request_cost({**deployment_pricing, "metadata": {"model_info": {"id": deployment_id}}}, token_hours) + return _vertex_request_cost({**token_prices, "metadata": {"model_info": {"id": deployment_id}}}, token_hours) def _vertex_request_cost(litellm_params: Mapping[str, object], token_hours: float) -> float | None: @@ -5976,17 +5979,26 @@ def test_a_storage_only_override_keeps_the_underlying_token_prices(_storage_cost cost_map_rate: Final = litellm.get_model_info(f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}")[_STORAGE_RATE_KEY] assert isinstance(cost_map_rate, float) deployment_rate: Final = cost_map_rate * 3 - storage_only: Final[dict[str, float]] = {_STORAGE_RATE_KEY: deployment_rate} no_override: Final = _vertex_request_cost({}, token_hours=0.0) - without_storage: Final = _routed_request_cost(storage_only, token_hours=0.0) - with_storage: Final = _routed_request_cost(storage_only, token_hours=500.0) + without_storage: Final = _routed_request_cost({}, token_hours=0.0, storage_rate=deployment_rate) + with_storage: Final = _routed_request_cost({}, token_hours=500.0, storage_rate=deployment_rate) assert no_override is not None and no_override > 0 assert without_storage == pytest.approx(no_override) assert with_storage == pytest.approx(no_override + 500.0 * deployment_rate) +def test_a_deployment_storage_rate_stays_off_the_shared_backend_key(_storage_cost_map: None) -> None: + backend_key: Final = f"vertex_ai/{_STORAGE_DEPLOYMENT_MODEL}" + cost_map_rate: Final = litellm.get_model_info(backend_key)[_STORAGE_RATE_KEY] + assert isinstance(cost_map_rate, float) + + _routed_request_cost({}, token_hours=0.0, storage_rate=cost_map_rate * 3) + + assert litellm.get_model_info(backend_key)[_STORAGE_RATE_KEY] == cost_map_rate + + def test_custom_pricing_as_completion_cost_param(): from litellm import Choices, Message, ModelResponse from litellm.utils import Usage diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 4aed84d05c8..d9404b3922e 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -1471,13 +1471,17 @@ def test_custom_pricing_field_denylist_covers_all_builtin_pricing_fields(): """ import typing - from litellm.types.utils import CustomPricingLiteLLMParams, ModelInfoBase + from litellm.types.utils import ( + DEPLOYMENT_SCOPED_SEPARATE_CHARGE_FIELDS, + CustomPricingLiteLLMParams, + ModelInfoBase, + ) pricing_markers = ("cost", "price", "uplift", "vector_size", "tiered_pricing") builtin_pricing_fields = { name for name in typing.get_type_hints(ModelInfoBase) if any(marker in name for marker in pricing_markers) } - denylisted_fields = set(CustomPricingLiteLLMParams.model_fields.keys()) + denylisted_fields = set(CustomPricingLiteLLMParams.model_fields.keys()) | DEPLOYMENT_SCOPED_SEPARATE_CHARGE_FIELDS uncovered = sorted(builtin_pricing_fields - denylisted_fields) assert not uncovered, ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ead1d315792..d91c0353323 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35967,8 +35967,6 @@ export interface components { cache_read_input_token_cost_priority?: number | null; /** Cache Read Input Token Cost Ultrafast */ cache_read_input_token_cost_ultrafast?: number | null; - /** Cache Storage Cost Per Token Per Hour */ - cache_storage_cost_per_token_per_hour?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Client Id */ @@ -50825,6 +50823,8 @@ export interface components { cache_creation_input_token_cost?: number | null; /** Cache Read Input Token Cost */ cache_read_input_token_cost?: number | null; + /** Cache Storage Cost Per Token Per Hour */ + cache_storage_cost_per_token_per_hour?: number | null; /** Cost Per Ptu Per Hour */ cost_per_ptu_per_hour?: number | null; /** Created At */ @@ -51166,8 +51166,6 @@ export interface components { cache_read_input_token_cost_priority?: number | null; /** Cache Read Input Token Cost Ultrafast */ cache_read_input_token_cost_ultrafast?: number | null; - /** Cache Storage Cost Per Token Per Hour */ - cache_storage_cost_per_token_per_hour?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Client Id */ From 0b07cb578893043b88b5918f2514c84f652e8113 Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:44:49 +0900 Subject: [PATCH 07/10] test(vertex_ai): cover cache storage joining provider additional costs --- tests/unit/test_cost_calculator.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 02570a38d79..2851ec8f766 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -1,4 +1,5 @@ import datetime +import math import time from collections.abc import Iterator, Mapping from pathlib import Path @@ -5927,6 +5928,30 @@ def test_context_cache_storage_uses_the_deployment_storage_rate(_storage_cost_ma assert with_storage - without_storage == pytest.approx(500.0 * deployment_rate) +def test_context_cache_storage_joins_the_provider_additional_costs(_storage_cost_map: None) -> None: + litellm.register_model( + {"router-storage-deployment": {"litellm_provider": "azure_ai", "mode": "chat", _STORAGE_RATE_KEY: 1e-6}} + ) + logging_obj: Final = _cache_storage_logging_obj("azure_ai/model-router", token_hours=500.0) + usage: Final = Usage(prompt_tokens=1_000_000, completion_tokens=100, total_tokens=1_000_100) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="azure_ai/gpt-4o-2024-08-06", usage=usage), + model="azure_ai/model-router", + custom_llm_provider="azure_ai", + litellm_logging_obj=logging_obj, + custom_pricing=True, + router_model_id="router-storage-deployment", + ) + + breakdown: Final = logging_obj.cost_breakdown + assert breakdown is not None + additional_costs: Final = breakdown.get("additional_costs") or {} + assert additional_costs["Azure Model Router Flat Cost"] > 0 + assert math.isclose(additional_costs["cache_storage_cost"], 500.0 * 1e-6) + assert math.isclose(cost, sum(additional_costs.values())) + + _STORAGE_DEPLOYMENT_MODEL: Final = "gemini-2.5-flash" From 77efb01bd8fdd04db68238ca47e3f402fc55d3d6 Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:58:33 +0900 Subject: [PATCH 08/10] fix(model-catalog): declare the cache_storage_cost_per_token_per_hour key The Rust model catalog files any key it does not know under problems, so the cost map entries that carry the Vertex/Gemini context cache storage rate fail the bundled registry check. Declare the key as a UnitPrice attribute, the same way thinking_on_by_default was declared. --- litellm-rust/crates/model-catalog/src/wire/attributes.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs index 401634f8c62..195d5bc9c0f 100644 --- a/litellm-rust/crates/model-catalog/src/wire/attributes.rs +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -53,6 +53,8 @@ attributes! { bedrock_converse_supports_strict_tools: bool, /// Highest reasoning effort the Bedrock output_config accepts for this model. bedrock_output_config_effort_ceiling: EffortLevel, + /// USD per token per hour to keep an explicitly created context cache stored. + cache_storage_cost_per_token_per_hour: UnitPrice, comment: String, /// Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'. default_reasoning_effort: EffortLevel, @@ -313,6 +315,8 @@ mod tests { #[case::deprecation_source_not_a_url_string("deprecation_source", json!(7), false)] #[case::thinking_on_by_default("thinking_on_by_default", json!(true), true)] #[case::thinking_on_by_default_not_a_flag("thinking_on_by_default", json!("yes"), false)] + #[case::cache_storage_rate("cache_storage_cost_per_token_per_hour", json!(4.5e-6), true)] + #[case::negative_cache_storage_rate("cache_storage_cost_per_token_per_hour", json!(-1), false)] #[case::reasoning_contexts("reasoning_contexts", json!(["reasoning", "non_reasoning"]), true)] #[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)] #[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)] From aa8bef28002e61fc906d12cba8f34cbb12a4ca7e Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:58:33 +0900 Subject: [PATCH 09/10] test(vertex_ai): cover storage recording on remembered and shared caches A cache found in the process-local known-cache record (sync and async) and the followers of a single-flight create reuse an existing cache, so they must record no storage token-hours; only the request that created the cache is billed for it. --- .../test_vertex_ai_context_caching.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) 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 a674764a40b..b6ccab05883 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,5 +1,7 @@ +import asyncio import json from collections.abc import Mapping +from datetime import datetime, timezone from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -2522,6 +2524,71 @@ async def test_async_reusing_an_existing_cache_records_no_storage(local_model_co assert reusing_logging_obj.context_cache_storage_token_hours == 0.0 +_STORAGE_TEST_CREATED_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp() + + +def _remembering_endpoints() -> ContextCachingEndpoints: + return ContextCachingEndpoints(known_caches=KnownCachedContents(clock=lambda: _STORAGE_TEST_CREATED_AT)) + + +def test_reusing_a_remembered_cache_records_no_storage(local_model_cost_map: None) -> None: + endpoints: Final = _remembering_endpoints() + client: Final = MagicMock(spec=HTTPHandler) + client.get.return_value = _json_response("GET", {}) + client.post.return_value = _json_response("POST", _CREATED_CACHE) + creating_logging_obj: Final = _storage_logging_obj() + endpoints.check_and_create_cache(**_create_kwargs(creating_logging_obj, client)) + reusing_logging_obj: Final = _storage_logging_obj() + + _, _, cache_name = endpoints.check_and_create_cache(**_create_kwargs(reusing_logging_obj, client)) + + assert cache_name == _STORAGE_TEST_CACHE_NAME + assert (client.get.call_count, client.post.call_count) == (1, 1) + assert creating_logging_obj.context_cache_storage_token_hours == 7200 * 0.5 + assert reusing_logging_obj.context_cache_storage_token_hours == 0.0 + + +@pytest.mark.asyncio +async def test_async_reusing_a_remembered_cache_records_no_storage(local_model_cost_map: None) -> None: + endpoints: Final = _remembering_endpoints() + client: Final = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=_json_response("GET", {})) + client.post = AsyncMock(return_value=_json_response("POST", _CREATED_CACHE)) + creating_logging_obj: Final = _storage_logging_obj() + await endpoints.async_check_and_create_cache(**_create_kwargs(creating_logging_obj, client)) + reusing_logging_obj: Final = _storage_logging_obj() + + _, _, cache_name = await endpoints.async_check_and_create_cache(**_create_kwargs(reusing_logging_obj, client)) + + assert cache_name == _STORAGE_TEST_CACHE_NAME + assert (client.get.await_count, client.post.await_count) == (1, 1) + assert creating_logging_obj.context_cache_storage_token_hours == 7200 * 0.5 + assert reusing_logging_obj.context_cache_storage_token_hours == 0.0 + + +@pytest.mark.asyncio +async def test_concurrent_requests_for_one_prefix_record_storage_once(local_model_cost_map: None) -> None: + endpoints: Final = _remembering_endpoints() + client: Final = MagicMock(spec=AsyncHTTPHandler) + + async def slow_list(**_: object) -> httpx.Response: + await asyncio.sleep(0.01) + return _json_response("GET", {}) + + client.get = AsyncMock(side_effect=slow_list) + client.post = AsyncMock(return_value=_json_response("POST", _CREATED_CACHE)) + logging_objs: Final = [_storage_logging_obj() for _ in range(3)] + + results: Final = await asyncio.gather( + *(endpoints.async_check_and_create_cache(**_create_kwargs(logging_obj, client)) for logging_obj in logging_objs) + ) + + assert {cache_name for _, _, cache_name in results} == {_STORAGE_TEST_CACHE_NAME} + assert client.post.await_count == 1 + recorded: Final = sorted(logging_obj.context_cache_storage_token_hours for logging_obj in logging_objs) + assert recorded == [0.0, 0.0, 7200 * 0.5] + + def _created_cache_body(total_token_count: int) -> dict: return { "name": "cachedContents/created", From 04833e4afb812d52d1eb2c83032979a69b586afd Mon Sep 17 00:00:00 2001 From: roasted-penguin <43606521+roasted-penguin@users.noreply.github.com> Date: Thu, 5 Nov 2026 12:31:47 +0900 Subject: [PATCH 10/10] fix(model-catalog): mirror cache_storage_cost_per_token_per_hour in ModelInfoFields ModelInfoFields lists the fields Python's ModelInfoBase(...) copies, and get_model_info_helper now copies cache_storage_cost_per_token_per_hour right before off_peak_pricing, so the mirror gets the same field in the same place. --- litellm-rust/crates/model-catalog/src/model_info.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index afbbfec59e6..2e0ca7805d2 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -58,6 +58,7 @@ pub struct ModelInfoFields { pub cache_creation_input_token_cost_above_12hr_ultrafast: Option, pub cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches: Option, pub cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast: Option, + pub cache_storage_cost_per_token_per_hour: Option, pub off_peak_pricing: Option, pub input_cost_per_character: Option, pub input_cost_per_token_above_128k_tokens: Option,