diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index 98a75748e..a510ed3bc 100644 --- a/google/genai/_extra_utils.py +++ b/google/genai/_extra_utils.py @@ -777,3 +777,27 @@ def get_usage_header( http_options.headers = existing_headers return config_model + + +def set_agent_platform_service_tier_header( + config: Optional[types.GenerateContentConfig], + is_vertex: bool, +) -> None: + """Injects the service_tier as an HTTP header for Gemini Enterprise Agent Platform requests.""" + if not is_vertex or not config or not config.service_tier: + return + + # Extract the raw string value (e.g. 'flex', 'standard', 'priority') + tier_val = ( + config.service_tier.value + if hasattr(config.service_tier, 'value') + else str(config.service_tier) + ) + tier_val = tier_val.lower().replace('service_tier_', '') + + if config.http_options is None: + config.http_options = types.HttpOptions(headers={}) + if config.http_options.headers is None: + config.http_options.headers = {} + + config.http_options.headers['X-Vertex-AI-LLM-Shared-Request-Type'] = tier_val diff --git a/google/genai/models.py b/google/genai/models.py index 45b20c724..3424f353f 100644 --- a/google/genai/models.py +++ b/google/genai/models.py @@ -1615,9 +1615,6 @@ def _GenerateContentConfig_to_vertex( getv(from_object, ['model_armor_config']), ) - if getv(from_object, ['service_tier']) is not None: - setv(parent_object, ['serviceTier'], getv(from_object, ['service_tier'])) - if getv(from_object, ['audio_transcription_config']) is not None: setv( to_object, @@ -6497,6 +6494,11 @@ def generate_content( _extra_utils.find_afc_incompatible_tool_indexes(config) ) parsed_config = _extra_utils.parse_config_for_mcp_usage(config) + + _extra_utils.set_agent_platform_service_tier_header( + parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False) + ) + if ( parsed_config and parsed_config.tools @@ -6663,6 +6665,11 @@ def generate_content_stream( _extra_utils.find_afc_incompatible_tool_indexes(config) ) parsed_config = _extra_utils.parse_config_for_mcp_usage(config) + + _extra_utils.set_agent_platform_service_tier_header( + parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False) + ) + if ( parsed_config and parsed_config.tools @@ -8628,6 +8635,10 @@ async def generate_content( else: parsed_config = config.model_copy(deep=True) + _extra_utils.set_agent_platform_service_tier_header( + parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False) + ) + # Use AsyncExitStack to keep MCP connections alive across the entire AFC loop async with contextlib.AsyncExitStack() as stack: @@ -8859,6 +8870,10 @@ async def generate_content_stream( else: parsed_config = config.model_copy(deep=True) + _extra_utils.set_agent_platform_service_tier_header( + parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False) + ) + incompatible_tools_indexes = ( _extra_utils.find_afc_incompatible_tool_indexes( parsed_config, diff --git a/google/genai/tests/models/test_generate_content.py b/google/genai/tests/models/test_generate_content.py index 38f35120a..2bde94138 100644 --- a/google/genai/tests/models/test_generate_content.py +++ b/google/genai/tests/models/test_generate_content.py @@ -577,24 +577,22 @@ class InstrumentEnum(Enum): pytest_helper.TestTableItem( name='test_service_tier', parameters=types._GenerateContentParameters( - model=GEMINI_FLASH_LATEST, + model='gemini-3.1-pro-preview', contents=t.t_contents('What is your name?'), config={ 'service_tier': 'FLEX', }, ), - exception_if_vertex='400', ), pytest_helper.TestTableItem( name='test_service_tier_lower', parameters=types._GenerateContentParameters( - model=GEMINI_FLASH_LATEST, + model='gemini-3.1-pro-preview', contents=t.t_contents('What is your name?'), config={ 'service_tier': 'flex', }, ), - exception_if_vertex='400', ), pytest_helper.TestTableItem( name='test_audio_transcription_config',