Skip to content

Commit 7726c5c

Browse files
authored
Add system_instruction to model repr (#272)
1 parent 7d2c17e commit 7726c5c

File tree

2 files changed

+14
-0
lines changed

2 files changed

+14
-0
lines changed

google/generativeai/generative_models.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,13 +103,19 @@ def model_name(self):
103103
return self._model_name
104104

105105
def __str__(self):
106+
def maybe_text(content):
107+
if content and len(content.parts) and (t := content.parts[0].text):
108+
return repr(t)
109+
return content
110+
106111
return textwrap.dedent(
107112
f"""\
108113
genai.GenerativeModel(
109114
model_name='{self.model_name}',
110115
generation_config={self._generation_config},
111116
safety_settings={self._safety_settings},
112117
tools={self._tools},
118+
system_instruction={maybe_text(self._system_instruction)},
113119
)"""
114120
)
115121

tests/test_generative_models.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1006,6 +1006,7 @@ def test_repr_for_multi_turn_chat(self):
10061006
generation_config={},
10071007
safety_settings={},
10081008
tools=None,
1009+
system_instruction=None,
10091010
),
10101011
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'})]
10111012
)"""
@@ -1033,6 +1034,7 @@ def test_repr_for_incomplete_streaming_chat(self):
10331034
generation_config={},
10341035
safety_settings={},
10351036
tools=None,
1037+
system_instruction=None,
10361038
),
10371039
history=[glm.Content({'parts': [{'text': 'I really like fantasy books.'}], 'role': 'user'}), <STREAMING IN PROGRESS>]
10381040
)"""
@@ -1076,12 +1078,18 @@ def test_repr_for_broken_streaming_chat(self):
10761078
generation_config={},
10771079
safety_settings={},
10781080
tools=None,
1081+
system_instruction=None,
10791082
),
10801083
history=[glm.Content({'parts': [{'text': 'I really like fantasy books.'}], 'role': 'user'}), <STREAMING ERROR>]
10811084
)"""
10821085
)
10831086
self.assertEqual(expected, result)
10841087

1088+
def test_repr_for_system_instruction(self):
1089+
model = generative_models.GenerativeModel("gemini-pro", system_instruction="Be excellent.")
1090+
result = repr(model)
1091+
self.assertIn("system_instruction='Be excellent.'", result)
1092+
10851093
def test_count_tokens_called_with_request_options(self):
10861094
self.client.count_tokens = unittest.mock.MagicMock()
10871095
request = unittest.mock.ANY

0 commit comments

Comments
 (0)