feat: Add ModelArmorConfig support for prompt and response sanitization via the Model Armor service

PiperOrigin-RevId: 857115428

Author: Abhi-13

google/genai/batches.py

@@ -1013,6 +1013,11 @@ def _GenerateContentConfig_to_mldev(
         getv(from_object, ['enable_enhanced_civic_answers']),
     )
 
+  if getv(from_object, ['model_armor_config']) is not None:
+    raise ValueError(
+        'model_armor_config parameter is not supported in Gemini API.'
+    )
+
   return to_object
 
 

google/genai/models.py

@@ -1117,6 +1117,11 @@ def _GenerateContentConfig_to_mldev(
         getv(from_object, ['enable_enhanced_civic_answers']),
     )
 
+  if getv(from_object, ['model_armor_config']) is not None:
+    raise ValueError(
+        'model_armor_config parameter is not supported in Gemini API.'
+    )
+
   return to_object
 
 
@@ -1279,6 +1284,13 @@ def _GenerateContentConfig_to_vertex(
         'enable_enhanced_civic_answers parameter is not supported in Vertex AI.'
     )
 
+  if getv(from_object, ['model_armor_config']) is not None:
+    setv(
+        parent_object,
+        ['modelArmorConfig'],
+        getv(from_object, ['model_armor_config']),
+    )
+
   return to_object
 
 

google/genai/tests/models/test_generate_content.py

@@ -546,6 +546,22 @@ class InstrumentEnum(Enum):
         ),
         exception_if_vertex='not supported',
     ),
+    pytest_helper.TestTableItem(
+        name='test_model_armor_config',
+        parameters=types._GenerateContentParameters(
+            model=GEMINI_FLASH_LATEST,
+            contents=t.t_contents('What is your name?'),
+            config={
+                'model_armor_config': {
+                    'prompt_template_name': '',
+                    'response_template_name': '',
+                    # Intentionally left blank just to test that the SDK doesn't
+                    # throw an exception.
+                },
+            },
+        ),
+        exception_if_mldev='not supported',
+    ),
 ]
 
 pytestmark = pytest_helper.setup(

google/genai/types.py

@@ -5003,6 +5003,38 @@ class SafetySettingDict(TypedDict, total=False):
 SpeechConfigUnionDict = Union[str, SpeechConfig, SpeechConfigDict]
 
 
+class ModelArmorConfig(_common.BaseModel):
+  """Configuration for Model Armor integrations of prompt and responses.
+
+  This data type is not supported in Gemini API.
+  """
+
+  prompt_template_name: Optional[str] = Field(
+      default=None,
+      description="""Optional. The name of the Model Armor template to use for prompt sanitization.""",
+  )
+  response_template_name: Optional[str] = Field(
+      default=None,
+      description="""Optional. The name of the Model Armor template to use for response sanitization.""",
+  )
+
+
+class ModelArmorConfigDict(TypedDict, total=False):
+  """Configuration for Model Armor integrations of prompt and responses.
+
+  This data type is not supported in Gemini API.
+  """
+
+  prompt_template_name: Optional[str]
+  """Optional. The name of the Model Armor template to use for prompt sanitization."""
+
+  response_template_name: Optional[str]
+  """Optional. The name of the Model Armor template to use for response sanitization."""
+
+
+ModelArmorConfigOrDict = Union[ModelArmorConfig, ModelArmorConfigDict]
+
+
 class GenerateContentConfig(_common.BaseModel):
   """Optional model configuration parameters.
 
@@ -5219,6 +5251,12 @@ class GenerateContentConfig(_common.BaseModel):
       models. This field is not supported in Vertex AI.
       """,
   )
+  model_armor_config: Optional[ModelArmorConfig] = Field(
+      default=None,
+      description="""Settings for prompt and response sanitization using the Model Armor
+      service. If supplied, safety_settings must not be supplied.
+      """,
+  )
 
   @pydantic.field_validator('response_schema', mode='before')
   @classmethod
@@ -5429,6 +5467,11 @@ class GenerateContentConfigDict(TypedDict, total=False):
       models. This field is not supported in Vertex AI.
       """
 
+  model_armor_config: Optional[ModelArmorConfigDict]
+  """Settings for prompt and response sanitization using the Model Armor
+      service. If supplied, safety_settings must not be supplied.
+      """
+
 
 GenerateContentConfigOrDict = Union[
     GenerateContentConfig, GenerateContentConfigDict