updated templates

Author: Sita04

model_armor/create_template.py

@@ -20,6 +20,8 @@ def create_model_armor_template(project_id: str, location: str, template_id: str
 
     from google.api_core.client_options import ClientOptions
     from google.cloud import modelarmor_v1
+    from google.cloud.modelarmor_v1 import Template, DetectionConfidenceLevel, FilterConfig, PiAndJailbreakFilterSettings, \
+        MaliciousUriFilterSettings
 
     client = modelarmor_v1.ModelArmorClient(
         transport="rest",
@@ -31,22 +33,18 @@ def create_model_armor_template(project_id: str, location: str, template_id: str
     # location = "us-central1"
     # template_id = "template_id"
 
-    template = {
-        "name": f"projects/{project_id}/locations/{location}/templates/{template_id}",
-        "filter_config": {
-            "rai_settings": {
-                "rai_filters": [
-                    {"filter_type": "HATE_SPEECH", "confidence_level": "LOW_AND_ABOVE"}
-                ]
-            },
-            "pi_and_jailbreak_filter_settings": {"filter_enforcement": "ENABLED"},
-            "malicious_uri_filter_settings": {"filter_enforcement": "ENABLED"},
-        },
-        "template_metadata": {
-            "log_template_operations": False,
-            "log_sanitize_operations": False,
-        },
-    }
+    template = Template(
+        name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
+        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 = modelarmor_v1.CreateTemplateRequest(

model_armor/sanitize_user_prompt.py

@@ -34,7 +34,7 @@ def sanitize_user_prompt(
     # template_id = "template_id"
 
     # Define the prompt
-    user_prompt = "My SSN is 123-45-6789"
+    user_prompt = "Can you describe this link? https://testsafebrowsing.appspot.com/s/malware.html"
 
     # Initialize request argument(s)
     user_prompt_data = modelarmor_v1.DataItem()

model_armor/test_templates.py

@@ -19,7 +19,6 @@
 from google.cloud.modelarmor_v1 import (
     DetectionConfidenceLevel,
     FilterMatchState,
-    RaiFilterType,
 )
 import pytest
 
@@ -60,10 +59,7 @@ def test_user_prompt() -> None:
 def test_update_templates() -> None:
     template = update_model_armor_template(PROJECT_ID, LOCATION, TEMPLATE_ID)
     assert (
-        template.filter_config.rai_settings.rai_filters[0].filter_type
-        == RaiFilterType.HATE_SPEECH
-        and template.filter_config.rai_settings.rai_filters[0].confidence_level
-        == DetectionConfidenceLevel.MEDIUM_AND_ABOVE
+        template.filter_config.pi_and_jailbreak_filter_settings.confidence_level == DetectionConfidenceLevel.LOW_AND_ABOVE
     )
 
 

model_armor/update_template.py

@@ -20,6 +20,8 @@ def update_model_armor_template(project_id: str, location: str, template_id: str
 
     from google.api_core.client_options import ClientOptions
     from google.cloud import modelarmor_v1
+    from google.cloud.modelarmor_v1 import Template, DetectionConfidenceLevel, FilterConfig, PiAndJailbreakFilterSettings, \
+        MaliciousUriFilterSettings
 
     client = modelarmor_v1.ModelArmorClient(
         transport="rest",
@@ -31,23 +33,18 @@ def update_model_armor_template(project_id: str, location: str, template_id: str
     # location = "us-central1"
     # template_id = "template_id"
 
-    updated_template = {
-        "name": f"projects/{project_id}/locations/{location}/templates/{template_id}",
-        "filter_config": {
-            "rai_settings": {
-                "rai_filters": [
-                    {
-                        "filter_type": "HATE_SPEECH",
-                        "confidence_level": "MEDIUM_AND_ABOVE",
-                    },
-                ]
-            },
-        },
-        "template_metadata": {
-            "log_template_operations": True,
-            "log_sanitize_operations": True,
-        },
-    }
+    updated_template = Template(
+        name=f"projects/{project_id}/locations/{location}/templates/{template_id}",
+        filter_config=FilterConfig(
+            pi_and_jailbreak_filter_settings=PiAndJailbreakFilterSettings(
+                filter_enforcement=PiAndJailbreakFilterSettings.PiAndJailbreakFilterEnforcement.ENABLED,
+                confidence_level=DetectionConfidenceLevel.LOW_AND_ABOVE,
+            ),
+            malicious_uri_filter_settings=MaliciousUriFilterSettings(
+                filter_enforcement=MaliciousUriFilterSettings.MaliciousUriFilterEnforcement.ENABLED,
+            )
+        ),
+    )
 
     # Initialize request argument(s)
     request = modelarmor_v1.UpdateTemplateRequest(template=updated_template)