fix: GenAI - Fixed GAPIC-renamed keys (keys with trailing underscore) in structures returned by the to_dict methods.

When calling `to_dict` on objects that contain schema (e.g. `GenerationConfig`), correct keys are now used (e.g. `type` instead of `type_`).

PiperOrigin-RevId: 696384806

Author: Ark-kun

tests/unit/vertexai/test_generative_models.py

@@ -178,47 +178,47 @@
     },
 }
 _RENAMING_EXPECTED_SCHEMA = {
-    "type_": "OBJECT",
+    "type": "OBJECT",
     "min_properties": "1",  # NB: int64 is converted to string
     "max_properties": "3",
     "properties": {
         "names": {
-            "type_": "ARRAY",
+            "type": "ARRAY",
             "min_items": "2",
             "max_items": "4",
             "items": {
-                "type_": "STRING",
+                "type": "STRING",
                 "min_length": "3",
                 "max_length": "5",
             },
         },
         "date": {
             "any_of": [
                 {
-                    "type_": "STRING",
-                    "format_": "date",
+                    "type": "STRING",
+                    "format": "date",
                 },
                 {
                     "any_of": [
                         {
-                            "type_": "INTEGER",
+                            "type": "INTEGER",
                             "minimum": 20241001,
                         },
                     ],
                 },
             ],
         },
         "ordered": {
-            "type_": "OBJECT",
+            "type": "OBJECT",
             "properties": {
-                "a": {"type_": "STRING"},
-                "b": {"type_": "INTEGER"},
+                "a": {"type": "STRING"},
+                "b": {"type": "INTEGER"},
                 "c": {
-                    "type_": "OBJECT",
+                    "type": "OBJECT",
                     "properties": {
-                        "x": {"type_": "STRING"},
-                        "y": {"type_": "NUMBER"},
-                        "z": {"type_": "INTEGER"},
+                        "x": {"type": "STRING"},
+                        "y": {"type": "NUMBER"},
+                        "z": {"type": "INTEGER"},
                     },
                     "property_ordering": ["z", "y", "x"],  # explicit order kept
                 },

vertexai/generative_models/_generative_models.py

@@ -15,7 +15,6 @@
 """Classes for working with generative models."""
 # pylint: disable=bad-continuation, line-too-long, protected-access
 
-from collections.abc import Mapping
 import copy
 import io
 import json
@@ -30,6 +29,7 @@
     Iterable,
     List,
     Literal,
+    Mapping,
     Optional,
     Sequence,
     Type,
@@ -2953,11 +2953,19 @@ def _append_gapic_part(
 
 def _proto_to_dict(message) -> Dict[str, Any]:
     """Converts a proto-plus protobuf message to a dictionary."""
-    return type(message).to_dict(
+    # The best way to convert proto to dict is not trivial.
+    # Ideally, we want original keys in snake_case.
+    # The preserving_proto_field_name flag controls key names, but states have issues:
+    # `False` leads to keys using camelCase instead of snake_case.
+    # `True` leads to keys using snake_case, but has renamed names like `type_`.
+    # We needs to fix this issue using _fix_renamed_proto_dict_keys_in_place.
+    result = type(message).to_dict(
         message,
         including_default_value_fields=False,
         use_integers_for_enums=False,
     )
+    _fix_renamed_proto_dict_keys_in_place(result)
+    return result
 
 
 def _dict_to_proto(message_type: Type[T], message_dict: Dict[str, Any]) -> T:
@@ -2974,6 +2982,21 @@ def _dict_to_pretty_string(d: dict) -> str:
     return json.dumps(d, indent=2)
 
 
+def _fix_renamed_proto_dict_keys_in_place(d: Mapping[str, Any]):
+    """Fixes proto dict keys in place."""
+    for key, value in list(d.items()):
+        if key.endswith("_"):
+            new_key = key.rstrip("_")
+            del d[key]
+            d[new_key] = value
+        if isinstance(value, Mapping):
+            _fix_renamed_proto_dict_keys_in_place(value)
+        if isinstance(value, Sequence) and not isinstance(value, str):
+            for item in value:
+                if isinstance(item, Mapping):
+                    _fix_renamed_proto_dict_keys_in_place(item)
+
+
 _FORMAT_TO_MIME_TYPE = {
     "png": "image/png",
     "jpeg": "image/jpeg",