Skip to content

Commit

Permalink
fix: GenAI - Fixed GAPIC-renamed keys (keys with trailing underscore)…
Browse files Browse the repository at this point in the history
… 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
  • Loading branch information
Ark-kun authored and copybara-github committed Nov 14, 2024
1 parent 9a30c31 commit 9d00424
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 15 deletions.
26 changes: 13 additions & 13 deletions tests/unit/vertexai/test_generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
},
Expand Down
27 changes: 25 additions & 2 deletions vertexai/generative_models/_generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -30,6 +29,7 @@
Iterable,
List,
Literal,
Mapping,
Optional,
Sequence,
Type,
Expand Down Expand Up @@ -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:
Expand All @@ -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",
Expand Down

0 comments on commit 9d00424

Please sign in to comment.