fix: GenAI SDK client(multimodal) - preserve existing metadata when creating from bigframes

PiperOrigin-RevId: 901268958

Author: cleop-google

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

@@ -210,6 +210,46 @@ def test_create_dataset_from_bigframes(client, is_replay_mode):
         )
 
 
+@pytest.mark.skipif(
+    sys.version_info < (3, 10), reason="bigframes requires python 3.10 or higher"
+)
+@pytest.mark.usefixtures("mock_bigquery_client", "mock_import_bigframes")
+def test_create_dataset_from_bigframes_preserves_other_metadata(client, is_replay_mode):
+    import bigframes.pandas
+
+    dataframe = pd.DataFrame(
+        {
+            "col1": ["col1"],
+            "col2": ["col2"],
+        }
+    )
+    if is_replay_mode:
+        bf_dataframe = mock.MagicMock()
+        bf_dataframe.to_gbq.return_value = "temp_table_id"
+    else:
+        bf_dataframe = bigframes.pandas.DataFrame(dataframe)
+
+    dataset = client.datasets.create_from_bigframes(
+        dataframe=bf_dataframe,
+        target_table_id=BIGQUERY_TABLE_NAME,
+        multimodal_dataset={
+            "display_name": "test-from-bigframes",
+            "metadata": {
+                "gemini_request_read_config": {
+                    "assembled_request_column_name": "test_column"
+                }
+            }
+        },
+    )
+
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-from-bigframes"
+    assert dataset.metadata.gemini_request_read_config.assembled_request_column_name == "test_column"
+    assert dataset.metadata.input_config.bigquery_source.uri == (
+        f"bq://{BIGQUERY_TABLE_NAME}"
+    )
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
@@ -371,3 +411,44 @@ async def test_create_dataset_from_bigframes_async(client, is_replay_mode):
         pd.testing.assert_frame_equal(
             rows.to_dataframe(), dataframe, check_index_type=False
         )
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 10), reason="bigframes requires python 3.10 or higher"
+)
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_bigquery_client", "mock_import_bigframes")
+async def test_create_dataset_from_bigframes_preserves_other_metadata_async(client, is_replay_mode):
+    import bigframes.pandas
+
+    dataframe = pd.DataFrame(
+        {
+            "col1": ["col1"],
+            "col2": ["col2"],
+        }
+    )
+    if is_replay_mode:
+        bf_dataframe = mock.MagicMock()
+        bf_dataframe.to_gbq.return_value = "temp_table_id"
+    else:
+        bf_dataframe = bigframes.pandas.DataFrame(dataframe)
+
+    dataset = await client.aio.datasets.create_from_bigframes(
+        dataframe=bf_dataframe,
+        target_table_id=BIGQUERY_TABLE_NAME,
+        multimodal_dataset={
+            "display_name": "test-from-bigframes",
+            "metadata": {
+                "gemini_request_read_config": {
+                    "assembled_request_column_name": "test_column"
+                }
+            }
+        },
+    )
+
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-from-bigframes"
+    assert dataset.metadata.gemini_request_read_config.assembled_request_column_name == "test_column"
+    assert dataset.metadata.input_config.bigquery_source.uri == (
+        f"bq://{BIGQUERY_TABLE_NAME}"
+    )

vertexai/_genai/_datasets_utils.py

@@ -34,17 +34,6 @@
 T = TypeVar("T", bound=BaseModel)
 
 
-def create_from_response(model_type: Type[T], response: dict[str, Any]) -> T:
-    """Creates a model from a response."""
-    model_field_names = model_type.model_fields.keys()
-    filtered_response = {}
-    for key, value in response.items():
-        snake_key = common.camel_to_snake(key)
-        if snake_key in model_field_names:
-            filtered_response[snake_key] = value
-    return model_type(**filtered_response)
-
-
 def validate_multimodal_dataset_bigquery_uri(
     multimodal_dataset: common.MultimodalDataset,
 ) -> None: