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 f4932508458..36cb9cd49b2 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 @@ -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 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 479f94130d4..ddc404cb8c7 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 @@ -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():