feat: GenAI SDK client(multimodal) - Allow passing dataset ID in addition to full resource name in dataset methods.

PiperOrigin-RevId: 899573600

Author: cleop-google

tests/unit/vertexai/genai/replays/test_get_multimodal_datasets.py

@@ -41,6 +41,15 @@ def test_get_dataset_from_public_method(client):
     assert dataset.display_name == "test-display-name"
 
 
+def test_get_dataset_by_id(client):
+    dataset = client.datasets.get_multimodal_dataset(
+        name="8810841321427173376",
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.name == DATASET
+    assert dataset.display_name == "test-display-name"
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
@@ -67,3 +76,13 @@ async def test_get_dataset_from_public_method_async(client):
     assert isinstance(dataset, types.MultimodalDataset)
     assert dataset.name == DATASET
     assert dataset.display_name == "test-display-name"
+
+
+@pytest.mark.asyncio
+async def test_get_dataset_by_id_async(client):
+    dataset = await client.aio.datasets.get_multimodal_dataset(
+        name="8810841321427173376",
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.name == DATASET
+    assert dataset.display_name == "test-display-name"

vertexai/_genai/_datasets_utils.py

@@ -262,3 +262,11 @@ async def save_dataframe_to_bigquery_async(
     )
     await asyncio.to_thread(copy_job.result)
     await asyncio.to_thread(bq_client.delete_table, temp_table_id)
+
+
+def resolve_dataset_name(resource_name_or_id: str, project: str, location: str) -> str:
+    """Resolves a dataset name or ID to a full resource name."""
+    resource_prefix = f"projects/{project}/locations/{location}/datasets/"
+    if not resource_name_or_id.startswith(resource_prefix):
+        return resource_prefix + resource_name_or_id
+    return resource_name_or_id

vertexai/_genai/datasets.py

@@ -1130,8 +1130,8 @@ def get_multimodal_dataset(
 
         Args:
           name:
-            Required. name of a multimodal dataset. The name should be in
-            the format of "projects/{project}/locations/{location}/datasets/{dataset}".
+            Required. A fully-qualified resource name or ID of the dataset.
+            Example: "projects/.../locations/.../datasets/123" or "123".
           config:
             Optional. A configuration for getting the multimodal dataset. If not
             provided, the default configuration will be used.
@@ -1145,6 +1145,10 @@ def get_multimodal_dataset(
         elif not config:
             config = types.VertexBaseConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not name.startswith(resource_prefix):
+            name = resource_prefix + name
+
         return self._get_multimodal_dataset(config=config, name=name)
 
     def delete_multimodal_dataset(
@@ -1172,6 +1176,10 @@ def delete_multimodal_dataset(
         elif not config:
             config = types.VertexBaseConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not name.startswith(resource_prefix):
+            name = resource_prefix + name
+
         return self._delete_multimodal_dataset(config=config, name=name)
 
     def assemble(
@@ -1207,6 +1215,10 @@ def assemble(
         elif not config:
             config = types.AssembleDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not name.startswith(resource_prefix):
+            name = resource_prefix + name
+
         operation = self._assemble_multimodal_dataset(
             name=name,
             gemini_request_read_config=gemini_request_read_config,
@@ -1255,6 +1267,10 @@ def assess_tuning_resources(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = self._assess_multimodal_dataset(
             name=dataset_name,
             tuning_resource_usage_assessment_config=types.TuningResourceUsageAssessmentConfig(
@@ -1316,6 +1332,10 @@ def assess_tuning_validity(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = self._assess_multimodal_dataset(
             name=dataset_name,
             tuning_validation_assessment_config=types.TuningValidationAssessmentConfig(
@@ -1376,6 +1396,10 @@ def assess_batch_prediction_resources(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = self._assess_multimodal_dataset(
             name=dataset_name,
             batch_prediction_resource_usage_assessment_config=types.BatchPredictionResourceUsageAssessmentConfig(
@@ -1435,6 +1459,10 @@ def assess_batch_prediction_validity(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = self._assess_multimodal_dataset(
             name=dataset_name,
             batch_prediction_validation_assessment_config=types.BatchPredictionValidationAssessmentConfig(
@@ -2352,21 +2380,25 @@ async def get_multimodal_dataset(
 
         Args:
           name:
-            Required. name of a multimodal dataset. The name should be in
-            the format of "projects/{project}/locations/{location}/datasets/{dataset}".
+            Required. A fully-qualified resource name or ID of the dataset.
+            Example: "projects/.../locations/.../datasets/123" or "123".
           config:
             Optional. A configuration for getting the multimodal dataset. If not
             provided, the default configuration will be used.
 
         Returns:
-          A types.MultimodalDataset object representing the updated multimodal
+          A types.MultimodalDataset object representing the retrieved multimodal
           dataset.
         """
         if isinstance(config, dict):
             config = types.VertexBaseConfig(**config)
         elif not config:
             config = types.VertexBaseConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not name.startswith(resource_prefix):
+            name = resource_prefix + name
+
         return await self._get_multimodal_dataset(config=config, name=name)
 
     async def delete_multimodal_dataset(
@@ -2394,6 +2426,10 @@ async def delete_multimodal_dataset(
         elif not config:
             config = types.VertexBaseConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not name.startswith(resource_prefix):
+            name = resource_prefix + name
+
         return await self._delete_multimodal_dataset(config=config, name=name)
 
     async def assemble(
@@ -2429,6 +2465,10 @@ async def assemble(
         elif not config:
             config = types.AssembleDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not name.startswith(resource_prefix):
+            name = resource_prefix + name
+
         operation = await self._assemble_multimodal_dataset(
             name=name,
             gemini_request_read_config=gemini_request_read_config,
@@ -2477,6 +2517,10 @@ async def assess_tuning_resources(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = await self._assess_multimodal_dataset(
             name=dataset_name,
             tuning_resource_usage_assessment_config=types.TuningResourceUsageAssessmentConfig(
@@ -2538,6 +2582,10 @@ async def assess_tuning_validity(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = await self._assess_multimodal_dataset(
             name=dataset_name,
             tuning_validation_assessment_config=types.TuningValidationAssessmentConfig(
@@ -2598,6 +2646,10 @@ async def assess_batch_prediction_resources(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = await self._assess_multimodal_dataset(
             name=dataset_name,
             batch_prediction_resource_usage_assessment_config=types.BatchPredictionResourceUsageAssessmentConfig(
@@ -2657,6 +2709,10 @@ async def assess_batch_prediction_validity(
         elif not config:
             config = types.AssessDatasetConfig()
 
+        resource_prefix = f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/"
+        if not dataset_name.startswith(resource_prefix):
+            dataset_name = resource_prefix + dataset_name
+
         operation = await self._assess_multimodal_dataset(
             name=dataset_name,
             batch_prediction_validation_assessment_config=types.BatchPredictionValidationAssessmentConfig(