scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix: return service_tier on final steamed chunk

Browse source
This commit is contained in:
Mark McDonald 2026-04-01 14:09:18 +08:00 • committed by Sameer Kankute
parent c2d2df2277
commit 4c923f6c86
No known key found for this signature in database
2 changed files with 14 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3168,12 +3168,12 @@ class ModelResponseIterator:
"provider_specific_fields", {}
)["traffic_type"] = traffic_type
service_tier = self.response_headers.get("x-gemini-service-tier")
if service_tier:
if service_tier.lower() == "standard":
setattr(model_response, "service_tier", "default")
else:
setattr(model_response, "service_tier", service_tier.lower())
service_tier = self.response_headers.get("x-gemini-service-tier")
if service_tier:
if service_tier.lower() == "standard":
setattr(model_response, "service_tier", "default")
else:
setattr(model_response, "service_tier", service_tier.lower())
setattr(model_response, "usage", usage) # type: ignore

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

@ -3520,9 +3520,16 @@ def test_vertex_ai_service_tier_streaming():
logging_obj=MagicMock(),
response_headers={"x-gemini-service-tier": "FLEX"},
)
# Undefined when usageMetadata is missing
result = iterator.chunk_parser(chunk)
assert result.service_tier == "flex"
# But definitely set when usageMetadata is present
chunk_with_usage = {
"candidates": [{"content": {"parts": [{"text": "hi"}]}}],
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}
}
result_with_usage = iterator.chunk_parser(chunk_with_usage)
assert result_with_usage.service_tier == "flex"
def test_vertex_ai_service_tier_non_streaming():