Update docstrings, fix jpeg support. (#108)

* Update docstrings, add png support.

* black

* Improve typechecking in docs

Author: MarkDaoust

docs/build_docs.py

@@ -24,6 +24,9 @@
 import pathlib
 import re
 import textwrap
+import typing
+
+typing.TYPE_CHECKING = True
 
 from absl import app
 from absl import flags

google/generativeai/__init__.py

@@ -82,6 +82,7 @@
 from google.generativeai.embedding import embed_content
 
 from google.generativeai.generative_models import GenerativeModel
+from google.generativeai.generative_models import ChatSession
 
 from google.generativeai.text import generate_text
 from google.generativeai.text import generate_embeddings

google/generativeai/generative_models.py

@@ -17,11 +17,11 @@
 from google.generativeai.types import generation_types
 from google.generativeai.types import safety_types
 
-_GENERATE_CONTENT_ASYNC_DOC = """The async version of `Model.generate_content`."""
+_GENERATE_CONTENT_ASYNC_DOC = """The async version of `GenerativeModel.generate_content`."""
 
 _GENERATE_CONTENT_DOC = """A multipurpose function to generate responses from the model.
 
-This `GenerativeModel.generate_content` method can handle multimodal input, and multiturn
+This `GenerativeModel.generate_content` method can handle multimodal input, and multi-turn
 conversations.
 
 >>> model = genai.GenerativeModel('models/gemini-pro')
@@ -289,6 +289,15 @@ def start_chat(
         *,
         history: Iterable[content_types.StrictContentType] | None = None,
     ) -> ChatSession:
+        """Returns a `genai.ChatSession` attached to this model.
+
+        >>> model = genai.GenerativeModel()
+        >>> chat = model.start_chat(history=[...])
+        >>> response = chat.send_message("Hello?")
+
+        Arguments:
+            history: An iterable of `glm.Content` objects, or equvalents to initialize the session.
+        """
         if self._generation_config.get("candidate_count", 1) > 1:
             raise ValueError("Can't chat with `candidate_count > 1`")
         return ChatSession(

google/generativeai/types/content_types.py

@@ -43,17 +43,23 @@
 ]
 
 
-def pil_to_png_bytes(img):
+def pil_to_blob(img):
     bytesio = io.BytesIO()
-    img.save(bytesio, format="PNG")
+    if isinstance(img, PIL.PngImagePlugin.PngImageFile):
+        img.save(bytesio, format="PNG")
+        mime_type = "image/png"
+    else:
+        img.save(bytesio, format="JPEG")
+        mime_type = "image/jpeg"
     bytesio.seek(0)
-    return bytesio.read()
+    data = bytesio.read()
+    return glm.Blob(mime_type=mime_type, data=data)
 
 
 def image_to_blob(image) -> glm.Blob:
     if PIL is not None:
         if isinstance(image, PIL.Image.Image):
-            return glm.Blob(mime_type="image/png", data=pil_to_png_bytes(image))
+            return pil_to_blob(image)
 
     if IPython is not None:
         if isinstance(image, IPython.display.Image):
@@ -71,7 +77,7 @@ def image_to_blob(image) -> glm.Blob:
             return glm.Blob(mime_type=mime_type, data=image.data)
 
     raise TypeError(
-        "Could not convert image. epected an `Image` type"
+        "Could not convert image. expected an `Image` type"
         "(`PIL.Image.Image` or `IPython.display.Image`).\n"
         f"Got a: {type(image)}\n"
         f"Value: {image}"

tests/test_content.py

@@ -13,26 +13,40 @@
 
 
 HERE = pathlib.Path(__file__).parent
-TEST_IMAGE_PATH = HERE / "test_img.png"
-TEST_IMAGE_URL = "https://storage.googleapis.com/generativeai-downloads/data/test_img.png"
-TEST_IMAGE_DATA = TEST_IMAGE_PATH.read_bytes()
+TEST_PNG_PATH = HERE / "test_img.png"
+TEST_PNG_URL = "https://storage.googleapis.com/generativeai-downloads/data/test_img.png"
+TEST_PNG_DATA = TEST_PNG_PATH.read_bytes()
+
+TEST_JPG_PATH = HERE / "test_img.jpg"
+TEST_JPG_URL = "https://storage.googleapis.com/generativeai-downloads/data/test_img.jpg"
+TEST_JPG_DATA = TEST_JPG_PATH.read_bytes()
 
 
 class UnitTests(parameterized.TestCase):
     @parameterized.named_parameters(
-        ["PIL", PIL.Image.open(TEST_IMAGE_PATH)],
-        ["IPython", IPython.display.Image(filename=TEST_IMAGE_PATH)],
+        ["PIL", PIL.Image.open(TEST_PNG_PATH)],
+        ["IPython", IPython.display.Image(filename=TEST_PNG_PATH)],
     )
-    def test_image_to_blob(self, image):
+    def test_png_to_blob(self, image):
         blob = content_types.image_to_blob(image)
         self.assertIsInstance(blob, glm.Blob)
         self.assertEqual(blob.mime_type, "image/png")
         self.assertStartsWith(blob.data, b"\x89PNG")
 
     @parameterized.named_parameters(
-        ["BlobDict", {"mime_type": "image/png", "data": TEST_IMAGE_DATA}],
-        ["glm.Blob", glm.Blob(mime_type="image/png", data=TEST_IMAGE_DATA)],
-        ["Image", IPython.display.Image(filename=TEST_IMAGE_PATH)],
+        ["PIL", PIL.Image.open(TEST_JPG_PATH)],
+        ["IPython", IPython.display.Image(filename=TEST_JPG_PATH)],
+    )
+    def test_jpg_to_blob(self, image):
+        blob = content_types.image_to_blob(image)
+        self.assertIsInstance(blob, glm.Blob)
+        self.assertEqual(blob.mime_type, "image/jpeg")
+        self.assertStartsWith(blob.data, b"\xff\xd8\xff\xe0\x00\x10JFIF")
+
+    @parameterized.named_parameters(
+        ["BlobDict", {"mime_type": "image/png", "data": TEST_PNG_DATA}],
+        ["glm.Blob", glm.Blob(mime_type="image/png", data=TEST_PNG_DATA)],
+        ["Image", IPython.display.Image(filename=TEST_PNG_PATH)],
     )
     def test_to_blob(self, example):
         blob = content_types.to_blob(example)
@@ -51,11 +65,11 @@ def test_to_part(self, example):
         self.assertEqual(part.text, "Hello world!")
 
     @parameterized.named_parameters(
-        ["Image", IPython.display.Image(filename=TEST_IMAGE_PATH)],
-        ["BlobDict", {"mime_type": "image/png", "data": TEST_IMAGE_DATA}],
+        ["Image", IPython.display.Image(filename=TEST_PNG_PATH)],
+        ["BlobDict", {"mime_type": "image/png", "data": TEST_PNG_DATA}],
         [
             "PartDict",
-            {"inline_data": {"mime_type": "image/png", "data": TEST_IMAGE_DATA}},
+            {"inline_data": {"mime_type": "image/png", "data": TEST_PNG_DATA}},
         ],
     )
     def test_img_to_part(self, example):
@@ -83,9 +97,9 @@ def test_to_content(self, example):
         self.assertEqual(part.text, "Hello world!")
 
     @parameterized.named_parameters(
-        ["ContentDict", {"parts": [PIL.Image.open(TEST_IMAGE_PATH)]}],
-        ["list[Image]", [PIL.Image.open(TEST_IMAGE_PATH)]],
-        ["Image", PIL.Image.open(TEST_IMAGE_PATH)],
+        ["ContentDict", {"parts": [PIL.Image.open(TEST_PNG_PATH)]}],
+        ["list[Image]", [PIL.Image.open(TEST_PNG_PATH)]],
+        ["Image", PIL.Image.open(TEST_PNG_PATH)],
     )
     def test_img_to_content(self, example):
         content = content_types.to_content(example)
@@ -140,10 +154,10 @@ def test_dict_to_content_fails(self):
     @parameterized.named_parameters(
         [
             "ContentDict",
-            [{"parts": [{"inline_data": PIL.Image.open(TEST_IMAGE_PATH)}]}],
+            [{"parts": [{"inline_data": PIL.Image.open(TEST_PNG_PATH)}]}],
         ],
-        ["ContentDict-unwraped", [{"parts": [PIL.Image.open(TEST_IMAGE_PATH)]}]],
-        ["Image", PIL.Image.open(TEST_IMAGE_PATH)],
+        ["ContentDict-unwraped", [{"parts": [PIL.Image.open(TEST_PNG_PATH)]}]],
+        ["Image", PIL.Image.open(TEST_PNG_PATH)],
     )
     def test_img_to_contents(self, example):
         contents = content_types.to_contents(example)