diff --git a/google/genai/batches.py b/google/genai/batches.py index 57ae78c49..5223a4df0 100644 --- a/google/genai/batches.py +++ b/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 diff --git a/google/genai/models.py b/google/genai/models.py index e81925fa0..0569ff3d8 100644 --- a/google/genai/models.py +++ b/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 diff --git a/google/genai/tests/models/test_generate_content.py b/google/genai/tests/models/test_generate_content.py index 60a1af172..bf5c0ea70 100644 --- a/google/genai/tests/models/test_generate_content.py +++ b/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( diff --git a/google/genai/types.py b/google/genai/types.py index a6190a069..e932e0373 100644 --- a/google/genai/types.py +++ b/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