diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 255c4a0902e..6992f9fb015 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2503,6 +2503,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata) setattr(model_response, "usage", usage) + if isinstance(cached_content_creation, dict): + model_response._hidden_params[VERTEX_AI_CACHED_CONTENT_KEY] = cached_content_creation ## ADD METADATA TO RESPONSE ## @@ -3265,6 +3267,9 @@ class ModelResponseIterator: VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata) + if isinstance(cached_content_creation, dict): + model_response._hidden_params[VERTEX_AI_CACHED_CONTENT_KEY] = cached_content_creation + traffic_type: Final = processed_chunk.get("usageMetadata", {}).get("trafficType") if traffic_type: model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index d31322c9ba4..35c827fff2a 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -6429,6 +6429,14 @@ def test_transform_response_applies_cache_creation_usage(include_creation): if include_creation: assert result.usage.prompt_tokens == 20010 assert result.usage.cache_creation_input_tokens == 10000 + assert result._hidden_params["vertex_ai_cached_content"] == { + "name": "cached-content", + "model": model, + "total_token_count": 10000, + "create_time": None, + "expire_time": None, + } else: assert result.usage.prompt_tokens == 10010 assert not hasattr(result.usage, "cache_creation_input_tokens") + assert "vertex_ai_cached_content" not in result._hidden_params