try to support enum types

Change-Id: I5141f751c4d6c578ef957aa8250cb26309ea9bd3

Author: MarkDaoust

google/generativeai/responder.py

@@ -116,7 +116,7 @@ def _generate_schema(
             inspect.Parameter.POSITIONAL_ONLY,
         )
     }
-    parameters = pydantic.create_model(f.__name__, **fields_dict).schema()
+    parameters = pydantic.create_model(f.__name__, **fields_dict).model_json_schema()
     # Postprocessing
     # 4. Suppress unnecessary title generation:
     #    * https://github.com/pydantic/pydantic/issues/1051

google/generativeai/types/content_types.py

@@ -379,7 +379,7 @@ def _schema_for_function(
 
 
 def _build_schema(fname, fields_dict):
-    parameters = pydantic.create_model(fname, **fields_dict).schema()
+    parameters = pydantic.create_model(fname, **fields_dict).model_json_schema()
     defs = parameters.pop("$defs", {})
     # flatten the defs
     for name, value in defs.items():
@@ -401,7 +401,10 @@ def _build_schema(fname, fields_dict):
 
 
 def unpack_defs(schema, defs):
-    properties = schema["properties"]
+    properties = schema.get("properties", None)
+    if properties is None:
+        return
+
     for name, value in properties.items():
         ref_key = value.get("$ref", None)
         if ref_key is not None:

tests/test_content.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import dataclasses
+import enum
 import pathlib
 import typing_extensions
 from typing import Any, Union, Iterable
@@ -64,6 +65,18 @@ class ADataClassWithNullable:
 class ADataClassWithList:
     a: list[int]
 
+class Choices(enum.Enum):
+    A:str = "A"
+    B:str = "B"
+    C:str = "C"
+    D:str = "D"
+
+
+@dataclasses.dataclass
+class HasEnum:
+    choice:Choices
+
+
 
 class UnitTests(parameterized.TestCase):
     @parameterized.named_parameters(
@@ -536,6 +549,20 @@ def b():
                 },
             ),
         ],
+        ['enum', Choices, protos.Schema(type=protos.Type.STRING, enum=['A', 'B', 'C', 'D'])],
+        ['enum_list',
+         list[Choices],
+         protos.Schema(
+              type="ARRAY",
+              items=protos.Schema(type=protos.Type.STRING, enum=['A', 'B', 'C', 'D']),
+         ),
+        ],
+        ['has_enum',
+         HasEnum,
+         protos.Schema(
+            type=protos.Type.OBJECT,
+            properties={'choice': protos.Schema(type=protos.Type.STRING, enum=['A', 'B', 'C', 'D'])})
+        ]
     )
     def test_auto_schema(self, annotation, expected):
         def fun(a: annotation):