feat: Added the ability to use the Model Armor service for content sanitization

Google APIs · copybara-github · commit 2d4fbff6d1f8 · 2025-07-30T10:40:47.000-07:00
PiperOrigin-RevId: 788957149
diff --git a/google/cloud/aiplatform/v1/content.proto b/google/cloud/aiplatform/v1/content.proto
@@ -17,6 +17,7 @@ syntax = "proto3";
 package google.cloud.aiplatform.v1;
 
 import "google/api/field_behavior.proto";
+import "google/api/resource.proto";
 import "google/cloud/aiplatform/v1/openapi.proto";
 import "google/cloud/aiplatform/v1/tool.proto";
 import "google/cloud/aiplatform/v1/vertex_rag_data.proto";
@@ -31,6 +32,10 @@ option java_outer_classname = "ContentProto";
 option java_package = "com.google.cloud.aiplatform.v1";
 option php_namespace = "Google\\Cloud\\AIPlatform\\V1";
 option ruby_package = "Google::Cloud::AIPlatform::V1";
+option (google.api.resource_definition) = {
+  type: "modelarmor.googleapis.com/Template"
+  pattern: "projects/{project}/locations/{location}/templates/{template}"
+};
 
 // Harm categories that will block the content.
 enum HarmCategory {
@@ -503,6 +508,9 @@ message Candidate {
 
     // The function call generated by the model is invalid.
     MALFORMED_FUNCTION_CALL = 9;
+
+    // The model response was blocked by Model Armor.
+    MODEL_ARMOR = 10;
   }
 
   // Output only. Index of the candidate.
@@ -725,6 +733,27 @@ message RetrievalMetadata {
       [(google.api.field_behavior) = OPTIONAL];
 }
 
+// Configuration for Model Armor integrations of prompt and responses.
+message ModelArmorConfig {
+  // Optional. The name of the Model Armor template to use for prompt
+  // sanitization.
+  string prompt_template_name = 1 [
+    (google.api.field_behavior) = OPTIONAL,
+    (google.api.resource_reference) = {
+      type: "modelarmor.googleapis.com/Template"
+    }
+  ];
+
+  // Optional. The name of the Model Armor template to use for response
+  // sanitization.
+  string response_template_name = 2 [
+    (google.api.field_behavior) = OPTIONAL,
+    (google.api.resource_reference) = {
+      type: "modelarmor.googleapis.com/Template"
+    }
+  ];
+}
+
 // Represents token counting info for a single modality.
 message ModalityTokenCount {
   // The modality associated with this token count.
diff --git a/google/cloud/aiplatform/v1/prediction_service.proto b/google/cloud/aiplatform/v1/prediction_service.proto
@@ -738,6 +738,11 @@ message GenerateContentRequest {
   repeated SafetySetting safety_settings = 3
       [(google.api.field_behavior) = OPTIONAL];
 
+  // Optional. Settings for prompt and response sanitization using the Model
+  // Armor service. If supplied, safety_settings must not be supplied.
+  ModelArmorConfig model_armor_config = 11
+      [(google.api.field_behavior) = OPTIONAL];
+
   // Optional. Generation config.
   GenerationConfig generation_config = 4
       [(google.api.field_behavior) = OPTIONAL];
@@ -764,6 +769,9 @@ message GenerateContentResponse {
 
       // Candidates blocked due to prohibited content.
       PROHIBITED_CONTENT = 4;
+
+      // The user prompt was blocked by Model Armor.
+      MODEL_ARMOR = 5;
     }
 
     // Output only. Blocked reason.
