scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

fix: use x-gemini-service-tier header in response

Browse source
This commit is contained in:
Mark McDonald 2026-03-26 13:33:55 +08:00 • committed by Sameer Kankute
parent 7617d4e6db
commit 75901ccbc6
No known key found for this signature in database
2 changed files with 21 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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