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
google-gemini · May 22, 2024 · 6df10a7 · 6df10a7
1 parent f361642
commit 6df10a7
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 107 deletions.
diff --git a/google/generativeai/generative_models.py b/google/generativeai/generative_models.py
@@ -3,10 +3,8 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-import dataclasses
 import textwrap
 from typing import Any
-from typing import Union
 import reprlib
 
 # pylint: disable=bad-continuation, line-too-long

diff --git a/google/generativeai/types/generation_types.py b/google/generativeai/types/generation_types.py
@@ -435,7 +435,9 @@ def __str__(self) -> str:
  else:
  _iterator = f"<{self._iterator.__class__.__name__}>"
 
- as_dict = self.to_dict()
+ as_dict = type(self._result).to_dict(
+ self._result, use_integers_for_enums=False, including_default_value_fields=False
+ )
  json_str = json.dumps(as_dict, indent=2)
 
  _result = f"glm.GenerateContentResponse({json_str})"

diff --git a/tests/test_generation.py b/tests/test_generation.py
@@ -531,13 +531,8 @@ def test_repr_for_generate_content_response_from_response(self):
  {
  "text": "Hello world!"
  }
- ],
- "role": ""
- },
- "finish_reason": 0,
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ ]
+ }
  }
  ]
  }),
@@ -567,13 +562,8 @@ def test_repr_for_generate_content_response_from_iterator(self):
  {
  "text": "a"
  }
- ],
- "role": ""
- },
- "finish_reason": 0,
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ ]
+ }
  }
  ]
  }),

diff --git a/tests/test_generative_models.py b/tests/test_generative_models.py
@@ -828,13 +828,8 @@ def test_repr_for_unary_non_streamed_response(self):
  {
  "text": "world!"
  }
- ],
- "role": ""
- },
- "finish_reason": 0,
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ ]
+ }
  }
  ]
  }),
@@ -866,13 +861,8 @@ def test_repr_for_streaming_start_to_finish(self):
  {
  "text": "first"
  }
- ],
- "role": ""
- },
- "finish_reason": 0,
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ ]
+ }
  }
  ]
  }),
@@ -896,28 +886,14 @@ def test_repr_for_streaming_start_to_finish(self):
  {
  "text": "first second"
  }
- ],
- "role": ""
+ ]
  },
  "index": 0,
- "citation_metadata": {
- "citation_sources": []
- },
- "finish_reason": 0,
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ "citation_metadata": {}
  }
  ],
- "prompt_feedback": {
- "block_reason": 0,
- "safety_ratings": []
- },
- "usage_metadata": {
- "prompt_token_count": 0,
- "candidates_token_count": 0,
- "total_token_count": 0
- }
+ "prompt_feedback": {},
+ "usage_metadata": {}
  }),
  )"""
  )
@@ -939,28 +915,14 @@ def test_repr_for_streaming_start_to_finish(self):
  {
  "text": "first second third"
  }
- ],
- "role": ""
+ ]
  },
  "index": 0,
- "citation_metadata": {
- "citation_sources": []
- },
- "finish_reason": 0,
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ "citation_metadata": {}
  }
  ],
- "prompt_feedback": {
- "block_reason": 0,
- "safety_ratings": []
- },
- "usage_metadata": {
- "prompt_token_count": 0,
- "candidates_token_count": 0,
- "total_token_count": 0
- }
+ "prompt_feedback": {},
+ "usage_metadata": {}
  }),
  )"""
  )
@@ -989,10 +951,8 @@ def test_repr_error_info_for_stream_prompt_feedback_blocked(self):
  iterator=<generator>,
  result=glm.GenerateContentResponse({
  "prompt_feedback": {
- "block_reason": 1,
- "safety_ratings": []
- },
- "candidates": []
+ "block_reason": "SAFETY"
+ }
  }),
  ),
  error=<BlockedPromptException> prompt_feedback {
@@ -1047,28 +1007,14 @@ def no_throw():
  {
  "text": "123"
  }
- ],
- "role": ""
+ ]
  },
  "index": 0,
- "citation_metadata": {
- "citation_sources": []
- },
- "finish_reason": 0,
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ "citation_metadata": {}
  }
  ],
- "prompt_feedback": {
- "block_reason": 0,
- "safety_ratings": []
- },
- "usage_metadata": {
- "prompt_token_count": 0,
- "candidates_token_count": 0,
- "total_token_count": 0
- }
+ "prompt_feedback": {},
+ "usage_metadata": {}
  }),
  ),
  error=<ValueError> """
@@ -1120,28 +1066,15 @@ def test_repr_error_info_for_chat_streaming_unexpected_stop(self):
  {
  "text": "abc"
  }
- ],
- "role": ""
+ ]
  },
- "finish_reason": 3,
+ "finish_reason": "SAFETY",
  "index": 0,
- "citation_metadata": {
- "citation_sources": []
- },
- "safety_ratings": [],
- "token_count": 0,
- "grounding_attributions": []
+ "citation_metadata": {}
  }
  ],
- "prompt_feedback": {
- "block_reason": 0,
- "safety_ratings": []
- },
- "usage_metadata": {
- "prompt_token_count": 0,
- "candidates_token_count": 0,
- "total_token_count": 0
- }
+ "prompt_feedback": {},
+ "usage_metadata": {}
  }),
  ),
  error=<StopCandidateException> index: 0