Keep required fields info in the structured-output schema

google-gemini · Nov 29, 2024 · bcb7cf9 · bcb7cf9
1 parent b8772ed
commit bcb7cf9
Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py
@@ -487,7 +487,6 @@ def strip_titles(schema):
 def add_object_type(schema):
     properties = schema.get("properties", None)
     if properties is not None:
-        schema.pop("required", None)
         schema["type"] = "object"
         for name, value in properties.items():
             add_object_type(value)

diff --git a/tests/test_content.py b/tests/test_content.py
@@ -557,6 +557,7 @@ def b():
             protos.Schema(
                 type=protos.Type.OBJECT,
                 properties={"a": {"type_": protos.Type.INTEGER}},
+                required=["a"],
             ),
         ],
         [
@@ -566,6 +567,7 @@ def b():
                 type=protos.Type.OBJECT,
                 nullable=True,
                 properties={"a": {"type_": protos.Type.INTEGER}},
+                required=["a"],
             ),
         ],
         [
@@ -576,6 +578,7 @@ def b():
                 items=protos.Schema(
                     type=protos.Type.OBJECT,
                     properties={"a": {"type_": protos.Type.INTEGER}},
+                    required="a",
                 ),
             ),
         ],
@@ -585,6 +588,7 @@ def b():
             protos.Schema(
                 type=protos.Type.OBJECT,
                 properties={"a": {"type_": protos.Type.INTEGER, "nullable": True}},
+                required=["a"],
             ),
         ],
         [
@@ -593,6 +597,7 @@ def b():
             protos.Schema(
                 type=protos.Type.OBJECT,
                 properties={"a": {"type_": "ARRAY", "items": {"type_": "INTEGER"}}},
+                required=["a"],
             ),
         ],
         [
@@ -602,6 +607,7 @@ def b():
                 items=protos.Schema(
                     type=protos.Type.OBJECT,
                     properties={"a": {"type_": "ARRAY", "items": {"type_": "INTEGER"}}},
+                    required="a",
                 ),
                 type="ARRAY",
             ),
@@ -622,6 +628,7 @@ def b():
                 properties={
                     "a": {"type_": protos.Type.INTEGER},
                 },
+                required=["a"],
             ),
         ],
         [
@@ -635,8 +642,10 @@ def b():
                         properties={
                             "a": {"type_": protos.Type.INTEGER},
                         },
+                        required=["a"],
                     ),
                 },
+                required="x",
             ),
         ],
         ["enum", Choices, protos.Schema(type=protos.Type.STRING, enum=["a", "b", "c", "d"])],
@@ -656,6 +665,7 @@ def b():
                 properties={
                     "choice": protos.Schema(type=protos.Type.STRING, enum=["a", "b", "c", "d"])
                 },
+                required=["choice"],
             ),
         ],
     )

diff --git a/tests/test_generation.py b/tests/test_generation.py
@@ -697,8 +697,10 @@ def test_repr_for_generate_content_response_from_iterator(self):
                             month=protos.Schema(type="INTEGER"),
                             year=protos.Schema(type="INTEGER"),
                         ),
+                        required=["day", "month", "year"],
                     ),
                 ),
+                required=["name", "favorite_color", "birthday"],
             ),
         ],
     )