chore: Automatically populating Schema.property_ordering in function …

…declaration schemas. PiperOrigin-RevId: 683138775
googleapis · Oct 7, 2024 · 66572e0 · 66572e0
1 parent 138dc1a
commit 66572e0
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 0 deletions.
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
@@ -158,6 +158,23 @@
                 },
             ],
         },
+        "ordered": {
+            "type": "OBJECT",
+            "properties": {
+                "a": {"type": "stRIng"},
+                "b": {"type": "Integer"},
+                "c": {
+                    "type": "objeCT",
+                    "properties": {
+                        "x": {"type": "string"},
+                        "y": {"type": "number"},
+                        "z": {"type": "integer"},
+                    },
+                    "property_ordering": ["z", "y", "x"],
+                },
+            },
+            "propertyOrdering": ["b", "a", "c"],
+        },
     },
 }
 _RENAMING_EXPECTED_SCHEMA = {
@@ -191,7 +208,25 @@
                 },
             ],
         },
+        "ordered": {
+            "type_": "OBJECT",
+            "properties": {
+                "a": {"type_": "STRING"},
+                "b": {"type_": "INTEGER"},
+                "c": {
+                    "type_": "OBJECT",
+                    "properties": {
+                        "x": {"type_": "STRING"},
+                        "y": {"type_": "NUMBER"},
+                        "z": {"type_": "INTEGER"},
+                    },
+                    "property_ordering": ["z", "y", "x"],  # explicit order kept
+                },
+            },
+            "property_ordering": ["b", "a", "c"],  # explicit order kept
+        },
     },
+    "property_ordering": ["names", "date", "ordered"],  # implicit order added
 }
 
 

diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -2094,6 +2094,11 @@ def _fix_schema_dict_for_gapic_in_place(schema_dict: Dict[str, Any]) -> None:
     if properties := schema_dict.get("properties"):
         for property_schema in properties.values():
             _fix_schema_dict_for_gapic_in_place(property_schema)
+        if (
+            "property_ordering" not in schema_dict
+            and "propertyOrdering" not in schema_dict
+        ):
+            schema_dict["property_ordering"] = list(properties.keys())
 
     if any_of := (schema_dict.get("any_of") or schema_dict.get("anyOf")):
         for any_of_schema in any_of:

diff --git a/vertexai/tokenization/_tokenizers.py b/vertexai/tokenization/_tokenizers.py
@@ -370,6 +370,8 @@ def _schema_traverse(self, schema: openapi.Schema) -> openapi.Schema:
         if "required" in schema:
             self._texts.extend(schema.required)
             counted_schema.required = schema.required
+        if "property_ordering" in schema:
+            counted_schema.property_ordering = schema.property_ordering
         if "items" in schema:
             counted_schema_items = self._schema_traverse(schema.items)
             counted_schema.items = counted_schema_items