Add system_instruction to model repr (#272)

google-gemini · Apr 8, 2024 · 7726c5c · 7726c5c
1 parent 7d2c17e
commit 7726c5c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/google/generativeai/generative_models.py b/google/generativeai/generative_models.py
@@ -103,13 +103,19 @@ def model_name(self):
         return self._model_name
 
     def __str__(self):
+        def maybe_text(content):
+            if content and len(content.parts) and (t := content.parts[0].text):
+                return repr(t)
+            return content
+
         return textwrap.dedent(
             f"""\
             genai.GenerativeModel(
                 model_name='{self.model_name}',
                 generation_config={self._generation_config},
                 safety_settings={self._safety_settings},
                 tools={self._tools},
+                system_instruction={maybe_text(self._system_instruction)},
             )"""
         )
 

diff --git a/tests/test_generative_models.py b/tests/test_generative_models.py
@@ -1006,6 +1006,7 @@ def test_repr_for_multi_turn_chat(self):
                     generation_config={},
                     safety_settings={},
                     tools=None,
+                    system_instruction=None,
                 ),
                 history=[glm.Content({'parts': [{'text': 'I really like fantasy books.'}], 'role': 'user'}), glm.Content({'parts': [{'text': 'first'}], 'role': 'model'}), glm.Content({'parts': [{'text': 'I also like this image.'}, {'inline_data': {'data': 'iVBORw0KGgoA...AAElFTkSuQmCC', 'mime_type': 'image/png'}}], 'role': 'user'}), glm.Content({'parts': [{'text': 'second'}], 'role': 'model'}), glm.Content({'parts': [{'text': 'What things do I like?.'}], 'role': 'user'}), glm.Content({'parts': [{'text': 'third'}], 'role': 'model'})]
             )"""
@@ -1033,6 +1034,7 @@ def test_repr_for_incomplete_streaming_chat(self):
                     generation_config={},
                     safety_settings={},
                     tools=None,
+                    system_instruction=None,
                 ),
                 history=[glm.Content({'parts': [{'text': 'I really like fantasy books.'}], 'role': 'user'}), <STREAMING IN PROGRESS>]
             )"""
@@ -1076,12 +1078,18 @@ def test_repr_for_broken_streaming_chat(self):
                     generation_config={},
                     safety_settings={},
                     tools=None,
+                    system_instruction=None,
                 ),
                 history=[glm.Content({'parts': [{'text': 'I really like fantasy books.'}], 'role': 'user'}), <STREAMING ERROR>]
             )"""
         )
         self.assertEqual(expected, result)
 
+    def test_repr_for_system_instruction(self):
+        model = generative_models.GenerativeModel("gemini-pro", system_instruction="Be excellent.")
+        result = repr(model)
+        self.assertIn("system_instruction='Be excellent.'", result)
+
     def test_count_tokens_called_with_request_options(self):
         self.client.count_tokens = unittest.mock.MagicMock()
         request = unittest.mock.ANY