diff --git a/google/generativeai/generative_models.py b/google/generativeai/generative_models.py index da981c860..02cab0b29 100644 --- a/google/generativeai/generative_models.py +++ b/google/generativeai/generative_models.py @@ -103,6 +103,11 @@ 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( @@ -110,6 +115,7 @@ def __str__(self): 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 index a7cdb3ff0..bc011823c 100644 --- 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'}), ] )""" @@ -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'}), ] )""" ) 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