feat: Add batch prediction assessments to multimodal dataset RPCs

Google APIs · copybara-github · commit f36c47534c2b · 2025-03-25T07:31:23.000-07:00
PiperOrigin-RevId: 740340685
diff --git a/google/cloud/aiplatform/v1beta1/dataset_service.proto b/google/cloud/aiplatform/v1beta1/dataset_service.proto
@@ -948,6 +948,18 @@ message AssessDataRequest {
     string model_name = 1 [(google.api.field_behavior) = REQUIRED];
   }
 
+  // Configuration for the batch prediction validation assessment.
+  message BatchPredictionValidationAssessmentConfig {
+    // Required. The name of the model used for batch prediction.
+    string model_name = 1 [(google.api.field_behavior) = REQUIRED];
+  }
+
+  // Configuration for the batch prediction resource usage assessment.
+  message BatchPredictionResourceUsageAssessmentConfig {
+    // Required. The name of the model used for batch prediction.
+    string model_name = 1 [(google.api.field_behavior) = REQUIRED];
+  }
+
   // The assessment type.
   oneof assessment_config {
     // Optional. Configuration for the tuning validation assessment.
@@ -958,6 +970,17 @@ message AssessDataRequest {
     TuningResourceUsageAssessmentConfig
         tuning_resource_usage_assessment_config = 3
         [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. Configuration for the batch prediction validation assessment.
+    BatchPredictionValidationAssessmentConfig
+        batch_prediction_validation_assessment_config = 6
+        [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. Configuration for the batch prediction resource usage
+    // assessment.
+    BatchPredictionResourceUsageAssessmentConfig
+        batch_prediction_resource_usage_assessment_config = 7
+        [(google.api.field_behavior) = OPTIONAL];
   }
 
   // The read config for the dataset.
@@ -966,6 +989,10 @@ message AssessDataRequest {
     // assess assembled data.
     GeminiTemplateConfig gemini_template_config = 4
         [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. The column name in the underlying table that contains already
+    // fully assembled requests.
+    string request_column_name = 5 [(google.api.field_behavior) = OPTIONAL];
   }
 
   // Required. The name of the Dataset resource. Used only for MULTIMODAL
@@ -997,6 +1024,18 @@ message AssessDataResponse {
     int64 billable_character_count = 2;
   }
 
+  // The result of the batch prediction validation assessment.
+  message BatchPredictionValidationAssessmentResult {}
+
+  // The result of the batch prediction resource usage assessment.
+  message BatchPredictionResourceUsageAssessmentResult {
+    // Number of tokens in the batch prediction dataset.
+    int64 token_count = 1;
+
+    // Number of audio tokens in the batch prediction dataset.
+    int64 audio_token_count = 2;
+  }
+
   // The assessment result.
   oneof assessment_result {
     // Optional. The result of the tuning validation assessment.
@@ -1007,6 +1046,16 @@ message AssessDataResponse {
     TuningResourceUsageAssessmentResult
         tuning_resource_usage_assessment_result = 2
         [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. The result of the batch prediction validation assessment.
+    BatchPredictionValidationAssessmentResult
+        batch_prediction_validation_assessment_result = 3
+        [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. The result of the batch prediction resource usage assessment.
+    BatchPredictionResourceUsageAssessmentResult
+        batch_prediction_resource_usage_assessment_result = 4
+        [(google.api.field_behavior) = OPTIONAL];
   }
 }
 
@@ -1110,6 +1159,11 @@ message AssembleDataRequest {
     // Optional. Config for assembling templates with a Gemini API structure.
     GeminiTemplateConfig gemini_template_config = 2
         [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. The column name in the underlying table that contains already
+    // fully assembled requests. If this field is set, the original request will
+    // be copied to the output table.
+    string request_column_name = 5 [(google.api.field_behavior) = OPTIONAL];
   }
 
   // Required. The name of the Dataset resource (used only for MULTIMODAL
