Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions google/genai/batches.py
Original file line number Diff line number Diff line change
Expand Up @@ -1052,6 +1052,9 @@ def _GenerateContentConfig_to_mldev(
'model_armor_config parameter is not supported in Gemini API.'
)

if getv(from_object, ['service_tier']) is not None:
setv(parent_object, ['serviceTier'], getv(from_object, ['service_tier']))

return to_object


Expand Down
6 changes: 6 additions & 0 deletions google/genai/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -1316,6 +1316,9 @@ def _GenerateContentConfig_to_mldev(
'model_armor_config parameter is not supported in Gemini API.'
)

if getv(from_object, ['service_tier']) is not None:
setv(parent_object, ['serviceTier'], getv(from_object, ['service_tier']))

return to_object


Expand Down Expand Up @@ -1502,6 +1505,9 @@ def _GenerateContentConfig_to_vertex(
getv(from_object, ['model_armor_config']),
)

if getv(from_object, ['service_tier']) is not None:
raise ValueError('service_tier parameter is not supported in Vertex AI.')

return to_object


Expand Down
33 changes: 20 additions & 13 deletions google/genai/tests/models/test_generate_content.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,9 +142,7 @@ class InstrumentEnum(Enum):
model=GEMINI_FLASH_LATEST,
contents=t.t_contents('high'),
config={
'system_instruction': t.t_content(
'I say high, you say low'
)
'system_instruction': t.t_content('I say high, you say low')
},
),
),
Expand Down Expand Up @@ -237,13 +235,13 @@ class InstrumentEnum(Enum):
config=types.GenerateContentConfig(
tools=[{'google_maps': {}}],
tool_config={
"retrieval_config": {
"lat_lng": {
"latitude": 37.421993,
"longitude": -122.079725,
'retrieval_config': {
'lat_lng': {
'latitude': 37.421993,
'longitude': -122.079725,
}
}
}
},
),
),
),
Expand Down Expand Up @@ -357,9 +355,7 @@ class InstrumentEnum(Enum):
name='test_speech_with_config',
parameters=types._GenerateContentParameters(
model='gemini-2.5-flash-preview-tts',
contents=t.t_contents(
'Produce a speech response saying "Cheese"'
),
contents=t.t_contents('Produce a speech response saying "Cheese"'),
config=types.GenerateContentConfig(
response_modalities=['audio'],
speech_config=types.SpeechConfig(
Expand Down Expand Up @@ -536,8 +532,8 @@ class InstrumentEnum(Enum):
parameters=types._GenerateContentParameters(
model=GEMINI_FLASH_LATEST,
contents=t.t_contents(
'Summarize the evidence that confirms widespread voter fraud was'
' the reason the last national election results were'
'Summarize the evidence that confirms widespread voter fraud'
' was the reason the last national election results were'
' inaccurate.'
),
config={
Expand All @@ -562,6 +558,17 @@ class InstrumentEnum(Enum):
),
exception_if_mldev='not supported',
),
pytest_helper.TestTableItem(
name='test_service_tier',
parameters=types._GenerateContentParameters(
model=GEMINI_FLASH_LATEST,
contents=t.t_contents('What is your name?'),
config={
'service_tier': 'flex',
},
),
exception_if_vertex='not supported',
),
]

pytestmark = pytest_helper.setup(
Expand Down
20 changes: 20 additions & 0 deletions google/genai/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -831,6 +831,19 @@ class ResourceScope(_common.CaseInSensitiveEnum):
"https://aiplatform.googleapis.com/publishers/google/models/gemini-3-pro-preview"""


class ServiceTier(_common.CaseInSensitiveEnum):
"""Pricing and performance service tier."""

SERVICE_TIER_UNSPECIFIED = 'SERVICE_TIER_UNSPECIFIED'
"""Default service tier, which is standard."""
SERVICE_TIER_FLEX = 'SERVICE_TIER_FLEX'
"""Flex service tier."""
SERVICE_TIER_STANDARD = 'SERVICE_TIER_STANDARD'
"""Standard service tier."""
SERVICE_TIER_PRIORITY = 'SERVICE_TIER_PRIORITY'
"""Priority service tier."""


class JSONSchemaType(Enum):
"""The type of the data supported by JSON Schema.

Expand Down Expand Up @@ -5937,6 +5950,10 @@ class GenerateContentConfig(_common.BaseModel):
service. If supplied, safety_settings must not be supplied.
""",
)
service_tier: Optional[ServiceTier] = Field(
default=None,
description="""The service tier to use for the request. For example, FLEX.""",
)

@pydantic.field_validator('response_schema', mode='before')
@classmethod
Expand Down Expand Up @@ -6151,6 +6168,9 @@ class GenerateContentConfigDict(TypedDict, total=False):
service. If supplied, safety_settings must not be supplied.
"""

service_tier: Optional[ServiceTier]
"""The service tier to use for the request. For example, FLEX."""


GenerateContentConfigOrDict = Union[
GenerateContentConfig, GenerateContentConfigDict
Expand Down
Loading