chore: GenAI - Improved schema canonization. Recursing into any_of …

…when fixing schema dict. PiperOrigin-RevId: 681934053
googleapis · Oct 3, 2024 · 2565a33 · 2565a33
1 parent c670eeb
commit 2565a33
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 9 deletions.
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
@@ -143,8 +143,20 @@
             },
         },
         "date": {
-            "type": "strinG",
-            "format": "date",
+            "any_of": [
+                {
+                    "type": "strinG",
+                    "format": "date",
+                },
+                {
+                    "anyOf": [
+                        {
+                            "type": "inTegEr",
+                            "minimum": 20241001,
+                        },
+                    ],
+                },
+            ],
         },
     },
 }
@@ -164,8 +176,20 @@
             },
         },
         "date": {
-            "type_": "STRING",
-            "format_": "date",
+            "any_of": [
+                {
+                    "type_": "STRING",
+                    "format_": "date",
+                },
+                {
+                    "any_of": [
+                        {
+                            "type_": "INTEGER",
+                            "minimum": 20241001,
+                        },
+                    ],
+                },
+            ],
         },
     },
 }

diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -2085,17 +2085,20 @@ def _convert_schema_dict_to_gapic(schema_dict: Dict[str, Any]) -> Dict[str, Any]
 
 def _fix_schema_dict_for_gapic_in_place(schema_dict: Dict[str, Any]) -> None:
     """Converts a JsonSchema to a dict that the Schema proto class accepts."""
-    schema_dict["type"] = schema_dict["type"].upper()
+    if "type" in schema_dict:
+        schema_dict["type"] = schema_dict["type"].upper()
 
-    items_schema = schema_dict.get("items")
-    if items_schema:
+    if items_schema := schema_dict.get("items"):
         _fix_schema_dict_for_gapic_in_place(items_schema)
 
-    properties = schema_dict.get("properties")
-    if properties:
+    if properties := schema_dict.get("properties"):
         for property_schema in properties.values():
             _fix_schema_dict_for_gapic_in_place(property_schema)
 
+    if any_of := (schema_dict.get("any_of") or schema_dict.get("anyOf")):
+        for any_of_schema in any_of:
+            _fix_schema_dict_for_gapic_in_place(any_of_schema)
+
 
 class CallableFunctionDeclaration(FunctionDeclaration):
     """A function declaration plus a function."""