Merge branch 'main' into text-error

Change-Id: Ia01fef6de0c0787461127c0112cbfe4e8c4cf8d8

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

@@ -73,11 +73,34 @@
 
 
 def pil_to_blob(img):
+    # When you load an image with PIL you get a subclass of PIL.Image
+    # The subclass knows what file type it was loaded from it has a `.format` class attribute
+    # and the `get_format_mimetype` method. Convert these back to the same file type.
+    #
+    # The base image class doesn't know its file type, it just knows its mode.
+    # RGBA converts to PNG easily, P[allet] converts to GIF, RGB to GIF.
+    # But for anything else I'm not going to bother mapping it out (for now) let's just convert to RGB and send it.
+    #
+    # References:
+    #   - file formats: https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html
+    #   - image modes: https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes
+
     bytesio = io.BytesIO()
-    if isinstance(img, PIL.PngImagePlugin.PngImageFile) or img.mode == "RGBA":
+
+    get_mime = getattr(img, "get_format_mimetype", None)
+    if get_mime is not None:
+        # If the image is created from a file, convert back to the same file type.
+        img.save(bytesio, format=img.format)
+        mime_type = img.get_format_mimetype()
+    elif img.mode == "RGBA":
         img.save(bytesio, format="PNG")
         mime_type = "image/png"
+    elif img.mode == "P":
+        img.save(bytesio, format="GIF")
+        mime_type = "image/gif"
     else:
+        if img.mode != "RGB":
+            img = img.convert("RGB")
         img.save(bytesio, format="JPEG")
         mime_type = "image/jpeg"
     bytesio.seek(0)
@@ -379,7 +402,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 +424,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:

google/generativeai/types/generation_types.py

@@ -16,7 +16,6 @@
 
 import collections
 import contextlib
-import sys
 from collections.abc import Iterable, AsyncIterable, Mapping
 import dataclasses
 import itertools
@@ -165,7 +164,7 @@ class GenerationConfig:
     top_p: float | None = None
     top_k: int | None = None
     response_mime_type: str | None = None
-    response_schema: protos.Schema | Mapping[str, Any] | None = None
+    response_schema: protos.Schema | Mapping[str, Any] | type | None = None
 
 
 GenerationConfigType = Union[protos.GenerationConfig, GenerationConfigDict, GenerationConfig]
@@ -186,7 +185,8 @@ def _normalize_schema(generation_config):
         if not str(response_schema).startswith("list["):
             raise ValueError(
                 f"Invalid input: Could not understand the type of '{response_schema}'. "
-                "Expected one of the following types: `int`, `float`, `str`, `bool`, `typing_extensions.TypedDict`, `dataclass`, or `list[...]`."
+                "Expected one of the following types: `int`, `float`, `str`, `bool`, `enum`, "
+                "`typing_extensions.TypedDict`, `dataclass` or `list[...]`."
             )
         response_schema = content_types._schema_for_class(response_schema)
 

samples/controlled_generation.py

@@ -11,9 +11,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from absl.testing import absltest
+import pathlib
 
 import google.generativeai as genai
 
+media = pathlib.Path(__file__).parents[1] / "third_party"
+
 
 class UnitTests(absltest.TestCase):
     def test_json_controlled_generation(self):
@@ -22,6 +25,7 @@ def test_json_controlled_generation(self):
 
         class Recipe(typing.TypedDict):
             recipe_name: str
+            ingredients: list[str]
 
         model = genai.GenerativeModel("gemini-1.5-pro-latest")
         result = model.generate_content(
@@ -36,14 +40,57 @@ class Recipe(typing.TypedDict):
     def test_json_no_schema(self):
         # [START json_no_schema]
         model = genai.GenerativeModel("gemini-1.5-pro-latest")
-        prompt = """List a few popular cookie recipes using this JSON schema:
+        prompt = """List a few popular cookie recipes in JSON format.
+
+        Use this JSON schema:
 
-        Recipe = {'recipe_name': str}
+        Recipe = {'recipe_name': str, 'ingredients': list[str]}
         Return: list[Recipe]"""
         result = model.generate_content(prompt)
         print(result)
         # [END json_no_schema]
 
+    def test_json_enum(self):
+        # [START json_enum]
+        import enum
+
+        class Choice(enum.Enum):
+            PERCUSSION = "Percussion"
+            STRING = "String"
+            WOODWIND = "Woodwind"
+            BRASS = "Brass"
+            KEYBOARD = "Keyboard"
+
+        model = genai.GenerativeModel("gemini-1.5-pro-latest")
+
+        organ = genai.upload_file(media / "organ.jpg")
+        result = model.generate_content(
+            ["What kind of instrument is this:", organ],
+            generation_config=genai.GenerationConfig(
+                response_mime_type="application/json", response_schema=Choice
+            ),
+        )
+        print(result)  # "Keyboard"
+        # [END json_enum]
+
+    def test_json_enum_raw(self):
+        # [START json_enum_raw]
+        model = genai.GenerativeModel("gemini-1.5-pro-latest")
+
+        organ = genai.upload_file(media / "organ.jpg")
+        result = model.generate_content(
+            ["What kind of instrument is this:", organ],
+            generation_config=genai.GenerationConfig(
+                response_mime_type="application/json",
+                response_schema={
+                    "type": "STRING",
+                    "enum": ["Percussion", "String", "Woodwind", "Brass", "Keyboard"],
+                },
+            ),
+        )
+        print(result)  # "Keyboard"
+        # [END json_enum_raw]
+
 
 if __name__ == "__main__":
     absltest.main()

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
@@ -35,6 +36,10 @@
 TEST_JPG_URL = "https://storage.googleapis.com/generativeai-downloads/data/test_img.jpg"
 TEST_JPG_DATA = TEST_JPG_PATH.read_bytes()
 
+TEST_GIF_PATH = HERE / "test_img.gif"
+TEST_GIF_URL = "https://storage.googleapis.com/generativeai-downloads/data/test_img.gif"
+TEST_GIF_DATA = TEST_GIF_PATH.read_bytes()
+
 
 # simple test function
 def datetime():
@@ -65,6 +70,18 @@ class ADataClassWithList:
     a: list[int]
 
 
+class Choices(enum.Enum):
+    A = "a"
+    B = "b"
+    C = "c"
+    D = "d"
+
+
+@dataclasses.dataclass
+class HasEnum:
+    choice: Choices
+
+
 class UnitTests(parameterized.TestCase):
     @parameterized.named_parameters(
         ["PIL", PIL.Image.open(TEST_PNG_PATH)],
@@ -88,6 +105,17 @@ def test_jpg_to_blob(self, image):
         self.assertEqual(blob.mime_type, "image/jpeg")
         self.assertStartsWith(blob.data, b"\xff\xd8\xff\xe0\x00\x10JFIF")
 
+    @parameterized.named_parameters(
+        ["PIL", PIL.Image.open(TEST_GIF_PATH)],
+        ["P", PIL.Image.fromarray(np.zeros([6, 6, 3], dtype=np.uint8)).convert("P")],
+        ["IPython", IPython.display.Image(filename=TEST_GIF_PATH)],
+    )
+    def test_gif_to_blob(self, image):
+        blob = content_types.image_to_blob(image)
+        self.assertIsInstance(blob, protos.Blob)
+        self.assertEqual(blob.mime_type, "image/gif")
+        self.assertStartsWith(blob.data, b"GIF87a")
+
     @parameterized.named_parameters(
         ["BlobDict", {"mime_type": "image/png", "data": TEST_PNG_DATA}],
         ["protos.Blob", protos.Blob(mime_type="image/png", data=TEST_PNG_DATA)],
@@ -536,6 +564,25 @@ 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):