Skip to content

Commit 6df10a7

Browse files
MarkDaoustMarkDaoust
authored
Make printing less verbose. (#340)
* Make printing less verbose. Change-Id: Ie37c1b75e57427f7256eb3c323a60a2947dcf6b1 * Use string values for enums when printing. Change-Id: I9b07bd5998445cc4ac59abcae8429fd7dacc824a * format Change-Id: Ibbdfca4be9370acd246721fd42629a10dc4ca612
1 parent f361642 commit 6df10a7

File tree

4 files changed

+30
-107
lines changed

4 files changed

+30
-107
lines changed

‎google/generativeai/generative_models.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,8 @@
33
from __future__ import annotations
44

55
from collections.abc import Iterable
6-
import dataclasses
76
import textwrap
87
from typing import Any
9-
from typing import Union
108
import reprlib
119

1210
# pylint: disable=bad-continuation, line-too-long

‎google/generativeai/types/generation_types.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -435,7 +435,9 @@ def __str__(self) -> str:
435435
else:
436436
_iterator = f"<{self._iterator.__class__.__name__}>"
437437

438-
as_dict = self.to_dict()
438+
as_dict = type(self._result).to_dict(
439+
self._result, use_integers_for_enums=False, including_default_value_fields=False
440+
)
439441
json_str = json.dumps(as_dict, indent=2)
440442

441443
_result = f"glm.GenerateContentResponse({json_str})"

‎tests/test_generation.py

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -531,13 +531,8 @@ def test_repr_for_generate_content_response_from_response(self):
531531
{
532532
"text": "Hello world!"
533533
}
534-
],
535-
"role": ""
536-
},
537-
"finish_reason": 0,
538-
"safety_ratings": [],
539-
"token_count": 0,
540-
"grounding_attributions": []
534+
]
535+
}
541536
}
542537
]
543538
}),
@@ -567,13 +562,8 @@ def test_repr_for_generate_content_response_from_iterator(self):
567562
{
568563
"text": "a"
569564
}
570-
],
571-
"role": ""
572-
},
573-
"finish_reason": 0,
574-
"safety_ratings": [],
575-
"token_count": 0,
576-
"grounding_attributions": []
565+
]
566+
}
577567
}
578568
]
579569
}),

‎tests/test_generative_models.py

Lines changed: 23 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -828,13 +828,8 @@ def test_repr_for_unary_non_streamed_response(self):
828828
{
829829
"text": "world!"
830830
}
831-
],
832-
"role": ""
833-
},
834-
"finish_reason": 0,
835-
"safety_ratings": [],
836-
"token_count": 0,
837-
"grounding_attributions": []
831+
]
832+
}
838833
}
839834
]
840835
}),
@@ -866,13 +861,8 @@ def test_repr_for_streaming_start_to_finish(self):
866861
{
867862
"text": "first"
868863
}
869-
],
870-
"role": ""
871-
},
872-
"finish_reason": 0,
873-
"safety_ratings": [],
874-
"token_count": 0,
875-
"grounding_attributions": []
864+
]
865+
}
876866
}
877867
]
878868
}),
@@ -896,28 +886,14 @@ def test_repr_for_streaming_start_to_finish(self):
896886
{
897887
"text": "first second"
898888
}
899-
],
900-
"role": ""
889+
]
901890
},
902891
"index": 0,
903-
"citation_metadata": {
904-
"citation_sources": []
905-
},
906-
"finish_reason": 0,
907-
"safety_ratings": [],
908-
"token_count": 0,
909-
"grounding_attributions": []
892+
"citation_metadata": {}
910893
}
911894
],
912-
"prompt_feedback": {
913-
"block_reason": 0,
914-
"safety_ratings": []
915-
},
916-
"usage_metadata": {
917-
"prompt_token_count": 0,
918-
"candidates_token_count": 0,
919-
"total_token_count": 0
920-
}
895+
"prompt_feedback": {},
896+
"usage_metadata": {}
921897
}),
922898
)"""
923899
)
@@ -939,28 +915,14 @@ def test_repr_for_streaming_start_to_finish(self):
939915
{
940916
"text": "first second third"
941917
}
942-
],
943-
"role": ""
918+
]
944919
},
945920
"index": 0,
946-
"citation_metadata": {
947-
"citation_sources": []
948-
},
949-
"finish_reason": 0,
950-
"safety_ratings": [],
951-
"token_count": 0,
952-
"grounding_attributions": []
921+
"citation_metadata": {}
953922
}
954923
],
955-
"prompt_feedback": {
956-
"block_reason": 0,
957-
"safety_ratings": []
958-
},
959-
"usage_metadata": {
960-
"prompt_token_count": 0,
961-
"candidates_token_count": 0,
962-
"total_token_count": 0
963-
}
924+
"prompt_feedback": {},
925+
"usage_metadata": {}
964926
}),
965927
)"""
966928
)
@@ -989,10 +951,8 @@ def test_repr_error_info_for_stream_prompt_feedback_blocked(self):
989951
iterator=<generator>,
990952
result=glm.GenerateContentResponse({
991953
"prompt_feedback": {
992-
"block_reason": 1,
993-
"safety_ratings": []
994-
},
995-
"candidates": []
954+
"block_reason": "SAFETY"
955+
}
996956
}),
997957
),
998958
error=<BlockedPromptException> prompt_feedback {
@@ -1047,28 +1007,14 @@ def no_throw():
10471007
{
10481008
"text": "123"
10491009
}
1050-
],
1051-
"role": ""
1010+
]
10521011
},
10531012
"index": 0,
1054-
"citation_metadata": {
1055-
"citation_sources": []
1056-
},
1057-
"finish_reason": 0,
1058-
"safety_ratings": [],
1059-
"token_count": 0,
1060-
"grounding_attributions": []
1013+
"citation_metadata": {}
10611014
}
10621015
],
1063-
"prompt_feedback": {
1064-
"block_reason": 0,
1065-
"safety_ratings": []
1066-
},
1067-
"usage_metadata": {
1068-
"prompt_token_count": 0,
1069-
"candidates_token_count": 0,
1070-
"total_token_count": 0
1071-
}
1016+
"prompt_feedback": {},
1017+
"usage_metadata": {}
10721018
}),
10731019
),
10741020
error=<ValueError> """
@@ -1120,28 +1066,15 @@ def test_repr_error_info_for_chat_streaming_unexpected_stop(self):
11201066
{
11211067
"text": "abc"
11221068
}
1123-
],
1124-
"role": ""
1069+
]
11251070
},
1126-
"finish_reason": 3,
1071+
"finish_reason": "SAFETY",
11271072
"index": 0,
1128-
"citation_metadata": {
1129-
"citation_sources": []
1130-
},
1131-
"safety_ratings": [],
1132-
"token_count": 0,
1133-
"grounding_attributions": []
1073+
"citation_metadata": {}
11341074
}
11351075
],
1136-
"prompt_feedback": {
1137-
"block_reason": 0,
1138-
"safety_ratings": []
1139-
},
1140-
"usage_metadata": {
1141-
"prompt_token_count": 0,
1142-
"candidates_token_count": 0,
1143-
"total_token_count": 0
1144-
}
1076+
"prompt_feedback": {},
1077+
"usage_metadata": {}
11451078
}),
11461079
),
11471080
error=<StopCandidateException> index: 0

0 commit comments

Comments
 (0)