Skip to content

Commit 3ec7238

Browse files
Sita04Sita04
committed
feat(modelarmor-samples): add model armor template and prompt samples
1 parent 09e274c commit 3ec7238

File tree

8 files changed

+323
-0
lines changed

8 files changed

+323
-0
lines changed

‎model_armor/create_template.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
def create_model_armor_template(project_id: str, location: str, template_id: str):
16+
# [START modelarmor_create_template]
17+
18+
from google.cloud import modelarmor_v1
19+
client = modelarmor_v1.ModelArmorClient(transport="rest", client_options={
20+
"api_endpoint": "modelarmor.us-central1.rep.googleapis.com"})
21+
22+
# TODO(Developer): Uncomment these variables and initialize
23+
# project_id = "your-google-cloud-project-id"
24+
# location = "us-central1"
25+
# template_id = "template_id"
26+
27+
template = {
28+
"name": f"projects/{project_id}/locations/{location}/templates/{template_id}",
29+
"filter_config": {
30+
"rai_settings": {
31+
"rai_filters": [
32+
{
33+
"filter_type": "HATE_SPEECH",
34+
"confidence_level": "LOW_AND_ABOVE"
35+
}
36+
]
37+
},
38+
"pi_and_jailbreak_filter_settings": {
39+
"filter_enforcement": "ENABLED"
40+
},
41+
"malicious_uri_filter_settings": {
42+
"filter_enforcement": "ENABLED"
43+
}
44+
},
45+
"template_metadata": {
46+
"log_template_operations": False,
47+
"log_sanitize_operations": False
48+
}
49+
}
50+
51+
# Initialize request arguments
52+
request = modelarmor_v1.CreateTemplateRequest(
53+
parent=f"projects/{project_id}/locations/{location}",
54+
template_id=template_id,
55+
template=template
56+
)
57+
58+
# Make the request
59+
response = client.create_template(request=request)
60+
61+
# Response
62+
return response
63+
# [END modelarmor_create_template]

‎model_armor/delete_template.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
def delete_model_armor_template(project_id: str, location: str, template_id: str):
16+
# [START modelarmor_delete_template]
17+
18+
from google.cloud import modelarmor_v1
19+
client = modelarmor_v1.ModelArmorClient(transport="rest", client_options={
20+
"api_endpoint": "modelarmor.us-central1.rep.googleapis.com"})
21+
22+
# TODO(Developer): Uncomment these variables and initialize
23+
# project_id = "your-google-cloud-project-id"
24+
# location = "us-central1"
25+
# template_id = "template_id"
26+
27+
request = modelarmor_v1.DeleteTemplateRequest(
28+
name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
29+
)
30+
31+
# Make the request
32+
client.delete_template(request=request)
33+
34+
# [END modelarmor_delete_template]

‎model_armor/get_template.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
def get_model_armor_template(project_id: str, location: str, template_id: str):
16+
# [START modelarmor_get_template]
17+
18+
from google.cloud import modelarmor_v1
19+
client = modelarmor_v1.ModelArmorClient(transport="rest", client_options={
20+
"api_endpoint": "modelarmor.us-central1.rep.googleapis.com"})
21+
22+
# TODO(Developer): Uncomment these variables and initialize
23+
# project_id = "your-google-cloud-project-id"
24+
# location = "us-central1"
25+
# template_id = "template_id"
26+
27+
# Initialize request arguments
28+
request = modelarmor_v1.GetTemplateRequest(
29+
name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
30+
)
31+
32+
# Make the request
33+
response = client.get_template(request=request)
34+
35+
# Handle the response
36+
return response
37+
# [END modelarmor_get_template]

‎model_armor/list_templates.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def list_model_armor_templates(project_id: str, location: str):
17+
# [START modelarmor_list_templates]
18+
from google.cloud import modelarmor_v1
19+
client = modelarmor_v1.ModelArmorClient(transport="rest", client_options={
20+
"api_endpoint": "modelarmor.us-central1.rep.googleapis.com"})
21+
22+
# TODO(Developer): Uncomment these variables and initialize
23+
# project_id = "your-google-cloud-project-id"
24+
# location = "us-central1"
25+
26+
# Initialize request argument(s)
27+
request = modelarmor_v1.ListTemplatesRequest(
28+
parent=f"projects/{project_id}/locations/{location}"
29+
)
30+
31+
# Make the request
32+
response = client.list_templates(request=request)
33+
for template in response:
34+
print(template.name)
35+
36+
# Handle the response
37+
return response
38+
# [END modelarmor_list_templates]

‎model_armor/requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
google-cloud-modelarmor

‎model_armor/sanitize_user_prompt.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from google.cloud.modelarmor_v1 import SanitizeUserPromptResponse
16+
17+
18+
def sanitize_user_prompt(project_id: str, location: str, template_id: str) -> SanitizeUserPromptResponse:
19+
# [START modelarmor_sanitize_user_prompt]
20+
21+
from google.cloud import modelarmor_v1
22+
client = modelarmor_v1.ModelArmorClient(transport="rest", client_options={
23+
"api_endpoint": "modelarmor.us-central1.rep.googleapis.com"})
24+
25+
# TODO(Developer): Uncomment these variables and initialize
26+
# project_id = "your-google-cloud-project-id"
27+
# location = "us-central1"
28+
# template_id = "template_id"
29+
30+
# Define the prompt
31+
user_prompt = "My SSN is 123-45-6789"
32+
33+
# Initialize request argument(s)
34+
user_prompt_data = modelarmor_v1.DataItem()
35+
user_prompt_data.text = user_prompt
36+
37+
request = modelarmor_v1.SanitizeUserPromptRequest(
38+
name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
39+
user_prompt_data=user_prompt_data,
40+
)
41+
42+
# Make the request
43+
response = client.sanitize_user_prompt(request=request)
44+
# Match state is TRUE when the prompt is caught by one of the safety policies in the template.
45+
print(response.sanitization_result.filter_match_state)
46+
47+
# Handle the response
48+
return response
49+
# [END modelarmor_sanitize_user_prompt]

‎model_armor/test_templates.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import os
2+
import uuid
3+
4+
import pytest
5+
from google.api_core.exceptions import NotFound
6+
from google.cloud.modelarmor_v1 import RaiFilterType, DetectionConfidenceLevel, FilterMatchState
7+
8+
from model_armor.create_template import create_model_armor_template
9+
from model_armor.delete_template import delete_model_armor_template
10+
from model_armor.get_template import get_model_armor_template
11+
from model_armor.list_templates import list_model_armor_templates
12+
from model_armor.sanitize_user_prompt import sanitize_user_prompt
13+
from model_armor.update_template import update_model_armor_template
14+
15+
PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
16+
LOCATION = "us-central1"
17+
TEMPLATE_ID = f"test-model-armor-{uuid.uuid4()}"
18+
19+
def test_create_template():
20+
template = create_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
21+
assert template is not None
22+
23+
def test_get_template():
24+
template = get_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
25+
assert TEMPLATE_ID in template.name
26+
27+
def test_list_templates():
28+
templates = list_model_armor_templates(PROJECT_ID, LOCATION)
29+
assert TEMPLATE_ID in str(templates)
30+
31+
def test_user_prompt():
32+
response = sanitize_user_prompt(PROJECT_ID, LOCATION, TEMPLATE_ID)
33+
assert response.sanitization_result.filter_match_state == FilterMatchState.MATCH_FOUND
34+
35+
def test_update_templates():
36+
template = update_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
37+
assert (template.filter_config.rai_settings.rai_filters[0].filter_type == RaiFilterType.HATE_SPEECH and
38+
template.filter_config.rai_settings.rai_filters[0].confidence_level == DetectionConfidenceLevel.MEDIUM_AND_ABOVE)
39+
40+
def test_delete_template():
41+
delete_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
42+
with pytest.raises(NotFound) as exception_info:
43+
get_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
44+
assert TEMPLATE_ID in str(exception_info.value)

‎model_armor/update_template.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def update_model_armor_template(project_id: str, location: str, template_id: str):
17+
# [START modelarmor_update_template]
18+
19+
from google.cloud import modelarmor_v1
20+
client = modelarmor_v1.ModelArmorClient(transport="rest", client_options={
21+
"api_endpoint": "modelarmor.us-central1.rep.googleapis.com"})
22+
23+
# TODO(Developer): Uncomment these variables and initialize
24+
# project_id = "your-google-cloud-project-id"
25+
# location = "us-central1"
26+
# template_id = "template_id"
27+
28+
updated_template = {
29+
"name": f"projects/{project_id}/locations/{location}/templates/{template_id}",
30+
"filter_config": {
31+
"rai_settings": {
32+
"rai_filters": [
33+
{
34+
"filter_type": "HATE_SPEECH",
35+
"confidence_level": "MEDIUM_AND_ABOVE"
36+
},
37+
]
38+
},
39+
},
40+
"template_metadata": {
41+
"log_template_operations": True,
42+
"log_sanitize_operations": True
43+
}
44+
45+
}
46+
47+
# Initialize request argument(s)
48+
request = modelarmor_v1.UpdateTemplateRequest(
49+
template=updated_template
50+
)
51+
52+
# Make the request
53+
response = client.update_template(request=request)
54+
55+
# Response
56+
return response
57+
# [END modelarmor_update_template]

0 commit comments

Comments
 (0)