Allow empty contents with count_tokens (#342)

Change-Id: Ic20e2f88427d2e4fbc97847cf5c2df1f80a9a5a1

Author: MarkDaoust

google/generativeai/generative_models.py

@@ -129,9 +129,6 @@ def _prepare_request(
         tool_config: content_types.ToolConfigType | None,
     ) -> glm.GenerateContentRequest:
         """Creates a `glm.GenerateContentRequest` from raw inputs."""
-        if not contents:
-            raise TypeError("contents must not be empty")
-
         tools_lib = self._get_tools_lib(tools)
         if tools_lib is not None:
             tools_lib = tools_lib.to_proto()
@@ -235,6 +232,9 @@ def generate_content(
             tools: `glm.Tools` more info coming soon.
             request_options: Options for the request.
         """
+        if not contents:
+            raise TypeError("contents must not be empty")
+
         request = self._prepare_request(
             contents=contents,
             generation_config=generation_config,
@@ -282,6 +282,9 @@ async def generate_content_async(
         request_options: helper_types.RequestOptionsType | None = None,
     ) -> generation_types.AsyncGenerateContentResponse:
         """The async version of `GenerativeModel.generate_content`."""
+        if not contents:
+            raise TypeError("contents must not be empty")
+
         request = self._prepare_request(
             contents=contents,
             generation_config=generation_config,

tests/test_generative_models.py

@@ -21,6 +21,10 @@
 TEST_IMAGE_DATA = TEST_IMAGE_PATH.read_bytes()
 
 
+def noop(x: int):
+    return x
+
+
 def simple_part(text: str) -> glm.Content:
     return glm.Content({"parts": [{"text": text}]})
 
@@ -725,18 +729,33 @@ def test_system_instruction(self, instruction, expected_instr):
         self.assertEqual(req.system_instruction, expected_instr)
 
     @parameterized.named_parameters(
-        ["basic", "Hello"],
-        ["list", ["Hello"]],
+        ["basic", {"contents": "Hello"}],
+        ["list", {"contents": ["Hello"]}],
         [
             "list2",
-            [{"text": "Hello"}, {"inline_data": {"data": b"PNG!", "mime_type": "image/png"}}],
+            {
+                "contents": [
+                    {"text": "Hello"},
+                    {"inline_data": {"data": b"PNG!", "mime_type": "image/png"}},
+                ]
+            },
         ],
-        ["contents", [{"role": "user", "parts": ["hello"]}]],
+        [
+            "contents",
+            {"contents": [{"role": "user", "parts": ["hello"]}]},
+        ],
+        ["empty", {}],
+        [
+            "system_instruction",
+            {"system_instruction": ["You are a cat"]},
+        ],
+        ["tools", {"tools": [noop]}],
     )
-    def test_count_tokens_smoke(self, contents):
+    def test_count_tokens_smoke(self, kwargs):
+        si = kwargs.pop("system_instruction", None)
         self.responses["count_tokens"] = [glm.CountTokensResponse(total_tokens=7)]
-        model = generative_models.GenerativeModel("gemini-pro-vision")
-        response = model.count_tokens(contents)
+        model = generative_models.GenerativeModel("gemini-pro-vision", system_instruction=si)
+        response = model.count_tokens(**kwargs)
         self.assertEqual(type(response).to_dict(response), {"total_tokens": 7})
 
     @parameterized.named_parameters(