scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(vertex_ai): surface cached content metadata on response hidden params

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-23 02:11:17 +00:00
parent 26833575a9
commit bcb6a69886
2 changed files with 13 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -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

8
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -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