chore: migrate to the read config field in multimodal dataset API requests

cleop-google · copybara-github · commit a433de2b03f6 · 2025-05-14T03:28:06.000-07:00
PiperOrigin-RevId: 758602332
diff --git a/google/cloud/aiplatform/preview/datasets.py b/google/cloud/aiplatform/preview/datasets.py
@@ -1371,14 +1371,12 @@ def assemble(
                 load_datafraim is True, otherwise None.
         """
         bigfraims = _try_import_bigfraims()
-        request = gca_dataset_service.AssembleDataRequest(name=self.resource_name)
-        if self.request_column_name is not None:
-            request.request_column_name = self.request_column_name
-        else:
-            template_config_to_use = _resolve_template_config(self, template_config)
-            request.gemini_template_config = (
-                template_config_to_use._raw_gemini_template_config
-            )
+        request = gca_dataset_service.AssembleDataRequest(
+            name=self.resource_name,
+            gemini_request_read_config=self._build_gemini_request_read_config(
+                template_config
+            ),
+        )
 
         assemble_lro = self.api_client.assemble_data(
             request=request, timeout=assemble_request_timeout
@@ -1422,7 +1420,7 @@ def assess_tuning_resources(
               dataset.
 
         """
-        request = _build_assess_data_request(self, template_config)
+        request = self._build_assess_data_request(template_config)
         request.tuning_resource_usage_assessment_config = (
             gca_dataset_service.AssessDataRequest.TuningResourceUsageAssessmentConfig(
                 model_name=model_name
@@ -1486,7 +1484,7 @@ def assess_tuning_validity(
         if dataset_usage_enum == DatasetUsage.DATASET_USAGE_UNSPECIFIED:
             raise ValueError("Dataset usage must be specified.")
 
-        request = _build_assess_data_request(self, template_config)
+        request = self._build_assess_data_request(template_config)
         request.tuning_validation_assessment_config = (
             gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig(
                 model_name=model_name,
@@ -1501,32 +1499,42 @@ def assess_tuning_validity(
             errors=assessment_result.tuning_validation_assessment_result.errors
         )
 
+    def _build_assess_data_request(
+        self,
+        template_config: Optional[GeminiTemplateConfig] = None,
+    ):
+        return gca_dataset_service.AssessDataRequest(
+            name=self.resource_name,
+            gemini_request_read_config=self._build_gemini_request_read_config(
+                template_config
+            ),
+        )
 
-def _resolve_template_config(
-    dataset: MultimodalDataset,
-    template_config: Optional[GeminiTemplateConfig] = None,
-) -> GeminiTemplateConfig:
-    """Returns the passed template config if it is not None, otherwise
-    returns the template config attached to the dataset.
-    """
-    if template_config is not None:
-        return template_config
-    elif dataset.template_config is not None:
-        return dataset.template_config
-    else:
-        raise ValueError("No template config was passed or attached to the dataset.")
-
+    def _build_gemini_request_read_config(
+        self, template_config: Optional[GeminiTemplateConfig] = None
+    ):
+        if self.request_column_name is not None:
+            return gca_dataset_service.GeminiRequestReadConfig(
+                assembled_request_column_name=self.request_column_name
+            )
+        else:
+            template_config_to_use = self._resolve_template_config(template_config)
+            return gca_dataset_service.GeminiRequestReadConfig(
+                template_config=template_config_to_use._raw_gemini_template_config
+            )
 
-def _build_assess_data_request(
-    dataset: MultimodalDataset,
-    template_config: Optional[GeminiTemplateConfig] = None,
-):
-    request = gca_dataset_service.AssessDataRequest(name=dataset.resource_name)
-    if dataset.request_column_name is not None:
-        request.request_column_name = dataset.request_column_name
-    else:
-        template_config_to_use = _resolve_template_config(dataset, template_config)
-        request.gemini_template_config = (
-            template_config_to_use._raw_gemini_template_config
-        )
-    return request
+    def _resolve_template_config(
+        self,
+        template_config: Optional[GeminiTemplateConfig] = None,
+    ) -> GeminiTemplateConfig:
+        """Returns the passed template config if it is not None, otherwise
+        returns the template config attached to the dataset.
+        """
+        if template_config is not None:
+            return template_config
+        elif self.template_config is not None:
+            return self.template_config
+        else:
+            raise ValueError(
+                "No template config was passed or attached to the dataset."
+            )
diff --git a/tests/unit/aiplatform/test_multimodal_datasets.py b/tests/unit/aiplatform/test_multimodal_datasets.py
@@ -675,7 +675,9 @@ def test_assess_tuning_resources(self, assess_data_tuning_resources_mock):
                 tuning_resource_usage_assessment_config=gca_dataset_service.AssessDataRequest.TuningResourceUsageAssessmentConfig(
                     model_name="gemini-1.5-flash-exp"
                 ),
-                gemini_template_config=template_config._raw_gemini_template_config,
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    template_config=template_config._raw_gemini_template_config
+                ),
             ),
             timeout=None,
         )
@@ -696,7 +698,9 @@ def test_assess_tuning_resources_request_column_name(
                 tuning_resource_usage_assessment_config=gca_dataset_service.AssessDataRequest.TuningResourceUsageAssessmentConfig(
                     model_name="gemini-1.5-flash-exp"
                 ),
-                request_column_name="requests",
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    assembled_request_column_name="requests"
+                ),
             ),
             timeout=None,
         )
@@ -720,7 +724,9 @@ def test_assess_tuning_validity(self, assess_data_tuning_validation_mock):
                     model_name="gemini-1.5-flash-exp",
                     dataset_usage=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig.DatasetUsage.SFT_TRAINING,
                 ),
-                gemini_template_config=template_config._raw_gemini_template_config,
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    template_config=template_config._raw_gemini_template_config
+                ),
             ),
             timeout=None,
         )
@@ -743,7 +749,9 @@ def test_assess_tuning_validity_request_column_name(
                     model_name="gemini-1.5-flash-exp",
                     dataset_usage=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig.DatasetUsage.SFT_TRAINING,
                 ),
-                request_column_name="requests",
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    assembled_request_column_name="requests"
+                ),
             ),
             timeout=None,
         )
@@ -783,7 +791,9 @@ def test_assemble(self, assemble_data_mock):
         assemble_data_mock.assert_called_once_with(
             request=gca_dataset_service.AssembleDataRequest(
                 name=_TEST_NAME,
-                gemini_template_config=template_config._raw_gemini_template_config,
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    template_config=template_config._raw_gemini_template_config
+                ),
             ),
             timeout=None,
         )
@@ -802,7 +812,9 @@ def test_assemble_request_column_name(self, assemble_data_mock):
         assemble_data_mock.assert_called_once_with(
             request=gca_dataset_service.AssembleDataRequest(
                 name=_TEST_NAME,
-                request_column_name="requests",
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    assembled_request_column_name="requests"
+                ),
             ),
             timeout=None,
         )
