feat: add support for model_selection_config to GenerateContentConfig

PiperOrigin-RevId: 745109883

Author: sararob

google/genai/models.py

@@ -176,6 +176,20 @@ def _Schema_to_mldev(
   return to_object
 
 
+def _ModelSelectionConfig_to_mldev(
+    api_client: BaseApiClient,
+    from_object: Union[dict, object],
+    parent_object: Optional[dict] = None,
+) -> dict:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['feature_selection_preference']) is not None:
+    raise ValueError(
+        'feature_selection_preference parameter is not supported in Gemini API.'
+    )
+
+  return to_object
+
+
 def _SafetySetting_to_mldev(
     api_client: BaseApiClient,
     from_object: Union[dict, object],
@@ -500,6 +514,11 @@ def _GenerateContentConfig_to_mldev(
   if getv(from_object, ['routing_config']) is not None:
     raise ValueError('routing_config parameter is not supported in Gemini API.')
 
+  if getv(from_object, ['model_selection_config']) is not None:
+    raise ValueError(
+        'model_selection_config parameter is not supported in Gemini API.'
+    )
+
   if getv(from_object, ['safety_settings']) is not None:
     setv(
         parent_object,
@@ -1273,6 +1292,22 @@ def _Schema_to_vertex(
   return to_object
 
 
+def _ModelSelectionConfig_to_vertex(
+    api_client: BaseApiClient,
+    from_object: Union[dict, object],
+    parent_object: Optional[dict] = None,
+) -> dict:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['feature_selection_preference']) is not None:
+    setv(
+        to_object,
+        ['featureSelectionPreference'],
+        getv(from_object, ['feature_selection_preference']),
+    )
+
+  return to_object
+
+
 def _SafetySetting_to_vertex(
     api_client: BaseApiClient,
     from_object: Union[dict, object],
@@ -1603,6 +1638,15 @@ def _GenerateContentConfig_to_vertex(
   if getv(from_object, ['routing_config']) is not None:
     setv(to_object, ['routingConfig'], getv(from_object, ['routing_config']))
 
+  if getv(from_object, ['model_selection_config']) is not None:
+    setv(
+        to_object,
+        ['modelConfig'],
+        _ModelSelectionConfig_to_vertex(
+            api_client, getv(from_object, ['model_selection_config']), to_object
+        ),
+    )
+
   if getv(from_object, ['safety_settings']) is not None:
     setv(
         parent_object,
@@ -2665,6 +2709,16 @@ def _GenerateVideosParameters_to_vertex(
   return to_object
 
 
+def _FeatureSelectionPreference_to_mldev_enum_validate(enum_value: Any):
+  if enum_value in set([
+      'FEATURE_SELECTION_PREFERENCE_UNSPECIFIED',
+      'PRIORITIZE_QUALITY',
+      'BALANCED',
+      'PRIORITIZE_COST',
+  ]):
+    raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
+
+
 def _SafetyFilterLevel_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set(['BLOCK_NONE']):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')

google/genai/tests/models/test_generate_content.py

@@ -441,6 +441,36 @@ def test_simple_config(client):
   assert response.text
 
 
+def test_model_selection_config_dict(client):
+  if not client.vertexai:
+    return
+  response = client.models.generate_content(
+      model='gemini-1.5-flash',
+      contents='Give me a Taylor Swift lyric and explain its meaning.',
+      config={
+          'model_selection_config': {
+              'feature_selection_preference': 'PRIORITIZE_COST'
+          }
+      },
+  )
+  assert response.text
+
+
+def test_model_selection_config_pydantic(client):
+  if not client.vertexai:
+    return
+  response = client.models.generate_content(
+      model='gemini-1.5-flash',
+      contents='Give me a Taylor Swift lyric and explain its meaning.',
+      config=types.GenerateContentConfig(
+          model_selection_config=types.ModelSelectionConfig(
+              feature_selection_preference=types.FeatureSelectionPreference.PRIORITIZE_QUALITY
+          )
+      ),
+  )
+  assert response.text
+
+
 def test_sdk_logger_logs_warnings(client, caplog):
   caplog.set_level(logging.DEBUG, logger='gemini_sdk_logger')
   sdk_logger = logging.getLogger('gemini_sdk_logger')

google/genai/types.py

@@ -236,6 +236,17 @@ class AdapterSize(_common.CaseInSensitiveEnum):
   ADAPTER_SIZE_THIRTY_TWO = 'ADAPTER_SIZE_THIRTY_TWO'
 
 
+class FeatureSelectionPreference(_common.CaseInSensitiveEnum):
+  """Options for feature selection preference."""
+
+  FEATURE_SELECTION_PREFERENCE_UNSPECIFIED = (
+      'FEATURE_SELECTION_PREFERENCE_UNSPECIFIED'
+  )
+  PRIORITIZE_QUALITY = 'PRIORITIZE_QUALITY'
+  BALANCED = 'BALANCED'
+  PRIORITIZE_COST = 'PRIORITIZE_COST'
+
+
 class DynamicRetrievalConfigMode(_common.CaseInSensitiveEnum):
   """Config for the dynamic retrieval config mode."""
 
@@ -1226,6 +1237,26 @@ class SchemaDict(TypedDict, total=False):
 SchemaOrDict = Union[Schema, SchemaDict]
 
 
+class ModelSelectionConfig(_common.BaseModel):
+  """Config for model selection."""
+
+  feature_selection_preference: Optional[FeatureSelectionPreference] = Field(
+      default=None, description="""Options for feature selection preference."""
+  )
+
+
+class ModelSelectionConfigDict(TypedDict, total=False):
+  """Config for model selection."""
+
+  feature_selection_preference: Optional[FeatureSelectionPreference]
+  """Options for feature selection preference."""
+
+
+ModelSelectionConfigOrDict = Union[
+    ModelSelectionConfig, ModelSelectionConfigDict
+]
+
+
 class SafetySetting(_common.BaseModel):
   """Safety settings."""
 
@@ -2387,6 +2418,11 @@ class GenerateContentConfig(_common.BaseModel):
       description="""Configuration for model router requests.
       """,
   )
+  model_selection_config: Optional[ModelSelectionConfig] = Field(
+      default=None,
+      description="""Configuration for model selection.
+      """,
+  )
   safety_settings: Optional[list[SafetySetting]] = Field(
       default=None,
       description="""Safety settings in the request to block unsafe content in the
@@ -2552,6 +2588,10 @@ class GenerateContentConfigDict(TypedDict, total=False):
   """Configuration for model router requests.
       """
 
+  model_selection_config: Optional[ModelSelectionConfigDict]
+  """Configuration for model selection.
+      """
+
   safety_settings: Optional[list[SafetySettingDict]]
   """Safety settings in the request to block unsafe content in the
       response.