chore: fix resolving between passed and attached read configs in multimodal datasets

PiperOrigin-RevId: 758655464

Author: cleop-google

google/cloud/aiplatform/preview/datasets.py

@@ -1511,30 +1511,24 @@ def _build_assess_data_request(
         )
 
     def _build_gemini_request_read_config(
-        self, template_config: Optional[GeminiTemplateConfig] = None
-    ):
-        if self.request_column_name is not None:
+        self, provided_template_config: Optional[GeminiTemplateConfig] = None
+    ) -> gca_dataset_service.GeminiRequestReadConfig:
+        """Returns the provided template config wrapped in a read config if it
+        is not None, otherwise returns the read config attached to the
+        dataset."""
+        if provided_template_config is not None:
             return gca_dataset_service.GeminiRequestReadConfig(
-                assembled_request_column_name=self.request_column_name
+                template_config=provided_template_config._raw_gemini_template_config
             )
-        else:
-            template_config_to_use = self._resolve_template_config(template_config)
+        elif self.template_config is not None:
             return gca_dataset_service.GeminiRequestReadConfig(
-                template_config=template_config_to_use._raw_gemini_template_config
+                template_config=self.template_config._raw_gemini_template_config
+            )
+        elif self.request_column_name is not None:
+            return gca_dataset_service.GeminiRequestReadConfig(
+                assembled_request_column_name=self.request_column_name
             )
-
-    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."
+                "No template config was provided and no read config is attached to the dataset."
             )

tests/unit/aiplatform/test_multimodal_datasets.py

@@ -116,6 +116,20 @@ def get_dataset_mock():
         yield get_dataset_mock
 
 
+@pytest.fixture
+def get_dataset_template_config_mock():
+    with mock.patch.object(
+        dataset_service.DatasetServiceClient, "get_dataset"
+    ) as get_dataset_mock:
+        get_dataset_mock.return_value = gca_dataset.Dataset(
+            display_name=_TEST_DISPLAY_NAME,
+            metadata_schema_uri=_TEST_METADATA_SCHEMA_URI_MULTIMODAL,
+            name=_TEST_NAME,
+            metadata=_TEST_METADATA_MULTIMODAL_WITH_TEMPLATE_CONFIG,
+        )
+        yield get_dataset_mock
+
+
 @pytest.fixture
 def get_dataset_request_column_name_mock():
     with mock.patch.object(
@@ -756,6 +770,60 @@ def test_assess_tuning_validity_request_column_name(
             timeout=None,
         )
 
+    @pytest.mark.usefixtures("get_dataset_template_config_mock")
+    def test_assess_tuning_validity_uses_attached_template_config(
+        self, assess_data_tuning_validation_mock
+    ):
+        aiplatform.init(project=_TEST_PROJECT)
+        dataset = ummd.MultimodalDataset(dataset_name=_TEST_NAME)
+        dataset.assess_tuning_validity(
+            model_name="gemini-1.5-flash-exp",
+            dataset_usage="SFT_TRAINING",
+        )
+        assess_data_tuning_validation_mock.assert_called_once_with(
+            request=gca_dataset_service.AssessDataRequest(
+                name=_TEST_NAME,
+                tuning_validation_assessment_config=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig(
+                    model_name="gemini-1.5-flash-exp",
+                    dataset_usage=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig.DatasetUsage.SFT_TRAINING,
+                ),
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    template_config=_TEST_METADATA_MULTIMODAL_WITH_TEMPLATE_CONFIG[
+                        "geminiTemplateConfigSource"
+                    ]["geminiTemplateConfig"]
+                ),
+            ),
+            timeout=None,
+        )
+
+    @pytest.mark.usefixtures("get_dataset_request_column_name_mock")
+    def test_assess_tuning_validity_request_column_name_overridden_by_template_config(
+        self, assess_data_tuning_validation_mock
+    ):
+        aiplatform.init(project=_TEST_PROJECT)
+        dataset = ummd.MultimodalDataset(dataset_name=_TEST_NAME)
+        template_config = ummd.GeminiTemplateConfig(
+            field_mapping={"question": "questionColumn"},
+        )
+        dataset.assess_tuning_validity(
+            model_name="gemini-1.5-flash-exp",
+            dataset_usage="SFT_TRAINING",
+            template_config=template_config,
+        )
+        assess_data_tuning_validation_mock.assert_called_once_with(
+            request=gca_dataset_service.AssessDataRequest(
+                name=_TEST_NAME,
+                tuning_validation_assessment_config=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig(
+                    model_name="gemini-1.5-flash-exp",
+                    dataset_usage=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig.DatasetUsage.SFT_TRAINING,
+                ),
+                gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
+                    template_config=template_config._raw_gemini_template_config
+                ),
+            ),
+            timeout=None,
+        )
+
     @pytest.mark.usefixtures("get_dataset_mock")
     def test_assess_tuning_validity_invalid_dataset_usage_throws_error(
         self, assess_data_tuning_validation_mock