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 bbb11394e1f..2b15b1862e0 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 @@ -2430,8 +2430,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): )["traffic_type"] = traffic_type ## ADD SERVICE TIER ## - if "serviceTier" in completion_response: - setattr(model_response, "service_tier", completion_response["serviceTier"]) + if getattr(raw_response, "headers", None): + if service_tier := raw_response.headers.get("x-gemini-service-tier"): + setattr(model_response, "service_tier", service_tier) except Exception as e: raise VertexAIError( @@ -2531,6 +2532,7 @@ async def make_call( streaming_response=response.aiter_lines(), sync_stream=False, logging_obj=logging_obj, + response_headers=response.headers, ) # LOGGING logging_obj.post_call( @@ -2573,6 +2575,7 @@ def make_sync_call( streaming_response=response.iter_lines(), sync_stream=True, logging_obj=logging_obj, + response_headers=response.headers, ) # LOGGING @@ -3029,7 +3032,11 @@ class VertexLLM(VertexBase): class ModelResponseIterator: def __init__( - self, streaming_response, sync_stream: bool, logging_obj: LoggingClass + self, + streaming_response, + sync_stream: bool, + logging_obj: LoggingClass, + response_headers: Optional[Dict[str, str]] = None, ): from litellm.litellm_core_utils.prompt_templates.common_utils import ( check_is_function_call, @@ -3040,6 +3047,7 @@ class ModelResponseIterator: self.accumulated_json = "" self.sent_first_chunk = False self.logging_obj = logging_obj + self.response_headers = response_headers or {} self.is_function_call = check_is_function_call(logging_obj) self.cumulative_tool_call_index: int = 0 self.has_seen_tool_calls: bool = False @@ -3157,8 +3165,9 @@ class ModelResponseIterator: "provider_specific_fields", {} )["traffic_type"] = traffic_type - if "serviceTier" in processed_chunk: - setattr(model_response, "service_tier", processed_chunk["serviceTier"]) + service_tier = self.response_headers.get("x-gemini-service-tier") + if service_tier: + setattr(model_response, "service_tier", service_tier) 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 7e1fff2fe2c..979785793da 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 @@ -3505,18 +3505,20 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming(): def test_vertex_ai_service_tier_streaming(): - """Test serviceTier is preserved in model_response for streaming.""" + """Test service_tier is preserved in model_response from headers for streaming.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, ) chunk = { "candidates": [{"content": {"parts": [{"text": "Hello"}]}}], - "serviceTier": "FLEX", } iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() + streaming_response=[], + sync_stream=True, + logging_obj=MagicMock(), + response_headers={"x-gemini-service-tier": "FLEX"}, ) result = iterator.chunk_parser(chunk) @@ -3524,7 +3526,7 @@ def test_vertex_ai_service_tier_streaming(): def test_vertex_ai_service_tier_non_streaming(): - """Test serviceTier is preserved in model_response for non-streaming.""" + """Test service_tier is preserved in model_response from headers for non-streaming.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -3541,11 +3543,11 @@ def test_vertex_ai_service_tier_non_streaming(): "candidatesTokenCount": 100, "totalTokenCount": 150, }, - "serviceTier": "FLEX", } raw_response = MagicMock() raw_response.json.return_value = completion_response + raw_response.headers = {"x-gemini-service-tier": "FLEX"} result = VertexGeminiConfig().transform_response( model="gemini-pro",