Fix response_schema with TypedDict

google-gemini · Nov 28, 2024 · c7e9884 · c7e9884
1 parent b8772ed
commit c7e9884
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py
@@ -423,7 +423,9 @@ def _build_schema(fname, fields_dict):
     for name, value in defs.items():
         unpack_defs(value, defs)
     unpack_defs(parameters, defs)
-
+    required_fields = []
+    if (fname == "dummy" and fields_dict["dummy"]):
+        required_fields = parameters['properties'][fname].get('required',None)
     # 5. Nullable fields:
     #     * https://github.com/pydantic/pydantic/issues/1270
     #     * https://stackoverflow.com/a/58841311
@@ -435,6 +437,8 @@ def _build_schema(fname, fields_dict):
     #    * https://github.com/pydantic/pydantic/issues/1051
     #    * http://cl/586221780
     strip_titles(parameters)
+    if required_fields:
+        parameters['properties']['dummy']['required']= required_fields
     return parameters
 
 

diff --git a/tests/test_generation.py b/tests/test_generation.py
@@ -699,6 +699,7 @@ def test_repr_for_generate_content_response_from_iterator(self):
                         ),
                     ),
                 ),
+                required=list(("name", "favorite_color", "birthday"))
             ),
         ],
     )
-Original file line number
+Diff line change
@@ Expand Up @@
                             ),
                         ),
                     ),
+                    required=list(("name", "favorite_color", "birthday"))
                 ),
             ],
         )
@@ Expand Down @@