feat(modelarmor-samples): add model armor template and prompt samples (#13143)

* feat(modelarmor-samples): add model armor template and prompt samples

* header fix

* add noxfile

* add pytest

* address comments and lint fix

* lint fix

* updated templates

* addressed review comments

* addressed review comments

* addressed review comments

Author: Sita04

model_armor/create_template.py

@@ -0,0 +1,68 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud.modelarmor_v1 import Template
+
+
+def create_model_armor_template(project_id: str, location: str, template_id: str) -> Template:
+    # [START modelarmor_create_template]
+
+    from google.api_core.client_options import ClientOptions
+    from google.cloud.modelarmor_v1 import (
+        Template,
+        DetectionConfidenceLevel,
+        FilterConfig,
+        PiAndJailbreakFilterSettings,
+        MaliciousUriFilterSettings,
+        ModelArmorClient,
+        CreateTemplateRequest
+    )
+
+    client = ModelArmorClient(
+        transport="rest",
+        client_options=ClientOptions(api_endpoint=f"modelarmor.{location}.rep.googleapis.com"),
+    )
+
+    # TODO(Developer): Uncomment these variables and initialize
+    # project_id = "your-google-cloud-project-id"
+    # location = "us-central1"
+    # template_id = "template_id"
+
+    template = Template(
+        filter_config=FilterConfig(
+            pi_and_jailbreak_filter_settings=PiAndJailbreakFilterSettings(
+                filter_enforcement=PiAndJailbreakFilterSettings.PiAndJailbreakFilterEnforcement.ENABLED,
+                confidence_level=DetectionConfidenceLevel.MEDIUM_AND_ABOVE,
+            ),
+            malicious_uri_filter_settings=MaliciousUriFilterSettings(
+                filter_enforcement=MaliciousUriFilterSettings.MaliciousUriFilterEnforcement.ENABLED,
+            )
+        ),
+    )
+
+    # Initialize request arguments
+    request = CreateTemplateRequest(
+        parent=f"projects/{project_id}/locations/{location}",
+        template_id=template_id,
+        template=template,
+    )
+
+    # Make the request
+    response = client.create_template(request=request)
+    # Response
+    print(response.name)
+
+# [END modelarmor_create_template]
+
+    return response

model_armor/delete_template.py

@@ -0,0 +1,43 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def delete_model_armor_template(project_id: str, location: str, template_id: str) -> None:
+    # [START modelarmor_delete_template]
+
+    from google.api_core.client_options import ClientOptions
+    from google.cloud.modelarmor_v1 import (
+        ModelArmorClient,
+        DeleteTemplateRequest,
+    )
+
+    client = ModelArmorClient(
+        transport="rest",
+        client_options=ClientOptions(api_endpoint=f"modelarmor.{location}.rep.googleapis.com"),
+    )
+
+    # TODO(Developer): Uncomment these variables and initialize
+    # project_id = "YOUR_PROJECT_ID"
+    # location = "us-central1"
+    # template_id = "template_id"
+
+    request = DeleteTemplateRequest(
+        name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
+    )
+
+    # Make the request
+    client.delete_template(request=request)
+
+
+# [END modelarmor_delete_template]

model_armor/get_template.py

@@ -0,0 +1,49 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud.modelarmor_v1 import Template
+
+
+def get_model_armor_template(project_id: str, location: str, template_id: str) -> Template:
+    # [START modelarmor_get_template]
+
+    from google.api_core.client_options import ClientOptions
+    from google.cloud.modelarmor_v1 import (
+        ModelArmorClient,
+        GetTemplateRequest,
+    )
+
+    client = ModelArmorClient(
+        transport="rest",
+        client_options=ClientOptions(api_endpoint=f"modelarmor.{location}.rep.googleapis.com"),
+    )
+
+    # TODO(Developer): Uncomment these variables and initialize
+    # project_id = "YOUR_PROJECT_ID"
+    # location = "us-central1"
+    # template_id = "template_id"
+
+    # Initialize request arguments
+    request = GetTemplateRequest(
+        name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
+    )
+
+    # Make the request
+    response = client.get_template(request=request)
+    print(response.name)
+
+# [END modelarmor_get_template]
+
+    # Handle the response
+    return response

model_armor/list_templates.py

@@ -0,0 +1,48 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud.modelarmor_v1.services.model_armor.pagers import ListTemplatesPager
+
+
+def list_model_armor_templates(project_id: str, location: str) -> ListTemplatesPager:
+    # [START modelarmor_list_templates]
+    from google.api_core.client_options import ClientOptions
+    from google.cloud.modelarmor_v1 import (
+        ModelArmorClient,
+        ListTemplatesRequest,
+    )
+
+    client = ModelArmorClient(
+        transport="rest",
+        client_options=ClientOptions(api_endpoint=f"modelarmor.{location}.rep.googleapis.com"),
+    )
+
+    # TODO(Developer): Uncomment these variables and initialize
+    # project_id = "YOUR_PROJECT_ID"
+    # location = "us-central1"
+
+    # Initialize request argument(s)
+    request = ListTemplatesRequest(
+        parent=f"projects/{project_id}/locations/{location}"
+    )
+
+    # Make the request
+    response = client.list_templates(request=request)
+    for template in response:
+        print(template.name)
+
+# [END modelarmor_list_templates]
+
+    # Handle the response
+    return response

model_armor/noxfile_config.py

@@ -0,0 +1,42 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be imported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    "ignored_versions": ["2.7", "3.7", "3.8", "3.10", "3.11", "3.12"],
+    # Old samples are opted out of enforcing Python type hints
+    # All new samples should feature them
+    "enforce_type_hints": True,
+    # An envvar key for determining the project id to use. Change it
+    # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+    # build specific Cloud project. You can also use your own string
+    # to use your own Cloud project.
+    "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+    # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
+    # If you need to use a specific version of pip,
+    # change pip_version_override to the string representation
+    # of the version number, for example, "20.2.4"
+    "pip_version_override": None,
+    # A dictionary you want to inject into your test. Don't put any
+    # secrets here. These values will override predefined values.
+    "envs": {},
+}

model_armor/requirements-test.txt

@@ -0,0 +1 @@
+pytest==8.3.4

model_armor/requirements.txt

@@ -0,0 +1 @@
+google-cloud-modelarmor==0.1.1

model_armor/sanitize_user_prompt.py

@@ -0,0 +1,61 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.cloud.modelarmor_v1 import SanitizeUserPromptResponse
+
+
+def sanitize_user_prompt(
+    project_id: str, location: str, template_id: str
+) -> SanitizeUserPromptResponse:
+    # [START modelarmor_sanitize_user_prompt]
+
+    from google.api_core.client_options import ClientOptions
+    from google.cloud.modelarmor_v1 import (
+        ModelArmorClient,
+        DataItem,
+        SanitizeUserPromptRequest
+    )
+
+    client = ModelArmorClient(
+        transport="rest",
+        client_options=ClientOptions(api_endpoint=f"modelarmor.{location}.rep.googleapis.com"),
+    )
+
+    # TODO(Developer): Uncomment these variables and initialize
+    # project_id = "YOUR_PROJECT_ID"
+    # location = "us-central1"
+    # template_id = "template_id"
+
+    # Define the prompt
+    user_prompt = "Can you describe this link? https://testsafebrowsing.appspot.com/s/malware.html"
+
+    # Initialize request argument(s)
+    user_prompt_data = DataItem(
+        text=user_prompt
+    )
+
+    request = SanitizeUserPromptRequest(
+        name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
+        user_prompt_data=user_prompt_data,
+    )
+
+    # Make the request
+    response = client.sanitize_user_prompt(request=request)
+    # Match state is TRUE when the prompt is caught by one of the safety policies in the template.
+    print(response.sanitization_result.filter_match_state)
+
+# [END modelarmor_sanitize_user_prompt]
+
+    # Handle the response
+    return response

model_armor/test_templates.py

@@ -0,0 +1,70 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import uuid
+
+from google.api_core.exceptions import NotFound
+from google.cloud.modelarmor_v1 import (
+    DetectionConfidenceLevel,
+    FilterMatchState,
+)
+import pytest
+
+from create_template import create_model_armor_template
+from delete_template import delete_model_armor_template
+from get_template import get_model_armor_template
+from list_templates import list_model_armor_templates
+from sanitize_user_prompt import sanitize_user_prompt
+from update_template import update_model_armor_template
+
+PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
+LOCATION = "us-central1"
+TEMPLATE_ID = f"test-model-armor-{uuid.uuid4()}"
+
+
+def test_create_template() -> None:
+    template = create_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
+    assert template is not None
+
+
+def test_get_template() -> None:
+    template = get_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
+    assert TEMPLATE_ID in template.name
+
+
+def test_list_templates() -> None:
+    templates = list_model_armor_templates(PROJECT_ID, LOCATION)
+    assert TEMPLATE_ID in str(templates)
+
+
+def test_user_prompt() -> None:
+    response = sanitize_user_prompt(PROJECT_ID, LOCATION, TEMPLATE_ID)
+    assert (
+        response.sanitization_result.filter_match_state == FilterMatchState.MATCH_FOUND
+    )
+
+
+def test_update_templates() -> None:
+    template = update_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
+    assert (
+        template.filter_config.pi_and_jailbreak_filter_settings.confidence_level == DetectionConfidenceLevel.LOW_AND_ABOVE
+    )
+
+
+def test_delete_template() -> None:
+    delete_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
+    with pytest.raises(NotFound) as exception_info:
+        get_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
+    assert TEMPLATE_ID in str(exception_info.value)