test generate_text safety settings, filters and feedback.

Author: MarkDaoust

google/generativeai/discuss.py

@@ -392,9 +392,9 @@ def __init__(self, **kwargs):
     @set_doc(discuss_types.ChatResponse.last.__doc__)
     def last(self) -> Optional[str]:
         if self.messages[-1]:
-          return self.messages[-1]["content"]
+            return self.messages[-1]["content"]
         else:
-          return None
+            return None
 
     @last.setter
     def last(self, message: discuss_types.MessageOptions):
@@ -410,9 +410,11 @@ def reply(
                 f"reply can't be called on an async client, use reply_async instead."
             )
         if self.last is None:
-            raise ValueError('The last response from the model did not return any candidates.\n'
-                             'Check the `.filters` attribute to see why the responses were filtered:\n'
-                             f'{self.filters}')
+            raise ValueError(
+                "The last response from the model did not return any candidates.\n"
+                "Check the `.filters` attribute to see why the responses were filtered:\n"
+                f"{self.filters}"
+            )
 
         request = self.to_dict()
         request.pop("candidates")
@@ -438,9 +440,6 @@ async def reply_async(
         request = _make_generate_message_request(**request)
         return await _generate_response_async(request=request, client=self._client)
 
-def _convert_filters_to_enums(filters):
-    for f in filters:
-        f['reason'] = safety_types.BlockedReason(f['reason'])
 
 def _build_chat_response(
     request: glm.GenerateMessageRequest,
@@ -456,7 +455,7 @@ def _build_chat_response(
     response = type(response).to_dict(response)
     response.pop("messages")
 
-    _convert_filters_to_enums(response['filters'])
+    safety_types.convert_filters_to_enums(response["filters"])
 
     if response["candidates"]:
         last = response["candidates"][0]

google/generativeai/docstring_utils.py

@@ -13,8 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 def strip_oneof(docstring):
     lines = docstring.splitlines()
     lines = [line for line in lines if ".. _oneof:" not in line]
     lines = [line for line in lines if "This field is a member of `oneof`_" not in line]
-    return "\n".join(lines)
+    return "\n".join(lines)

google/generativeai/text.py

@@ -45,7 +45,7 @@ def _make_generate_text_request(
     max_output_tokens: Optional[int] = None,
     top_p: Optional[int] = None,
     top_k: Optional[int] = None,
-    safety_settings: Optional[List[safety_types.SafetySetting]] = None,
+    safety_settings: Optional[List[safety_types.SafetySettingDict]] = None,
     stop_sequences: Union[str, Iterable[str]] = None,
 ) -> glm.GenerateTextRequest:
     model = model_types.make_model_name(model)
@@ -77,7 +77,7 @@ def generate_text(
     max_output_tokens: Optional[int] = None,
     top_p: Optional[float] = None,
     top_k: Optional[float] = None,
-    safety_settings: Optional[Iterable[safety.SafetySetting]] = None,
+    safety_settings: Optional[Iterable[safety.SafetySettingDict]] = None,
     stop_sequences: Union[str, Iterable[str]] = None,
     client: Optional[glm.TextServiceClient] = None,
 ) -> text_types.Completion:
@@ -159,6 +159,9 @@ def _generate_response(
     response = client.generate_text(request)
     response = type(response).to_dict(response)
 
+    safety_types.convert_filters_to_enums(response["filters"])
+    safety_types.convert_safety_feedback_to_enums(response["safety_feedback"])
+
     return Completion(_client=client, **response)
 
 

google/generativeai/types/__init__.py

@@ -24,4 +24,4 @@
 del model_types
 del text_types
 del citation_types
-del safety_types
+del safety_types

google/generativeai/types/discuss_types.py

@@ -162,7 +162,7 @@ class ChatResponse(abc.ABC):
     candidates: List[MessageDict]
     top_p: Optional[float] = None
     top_k: Optional[float] = None
-    filters: List[safety_types.ContentFilter]
+    filters: List[safety_types.ContentFilterDict]
 
     @property
     @abc.abstractmethod

google/generativeai/types/safety_types.py

@@ -23,9 +23,9 @@
     "HarmProbability",
     "HarmBlockThreshold",
     "BlockedReason",
-    "ContentFilter",
+    "ContentFilterDict",
     "SafetyRatingDict",
-    "SafetySetting",
+    "SafetySettingDict",
     "SafetyFeedbackDict",
 ]
 
@@ -36,29 +36,50 @@
 BlockedReason = glm.ContentFilter.BlockedReason
 
 
-class ContentFilter(TypedDict):
+class ContentFilterDict(TypedDict):
     reason: BlockedReason
     message: str
 
     __doc__ = docstring_utils.strip_oneof(glm.ContentFilter.__doc__)
 
 
+def convert_filters_to_enums(filters):
+    for f in filters:
+        f["reason"] = BlockedReason(f["reason"])
+
+
 class SafetyRatingDict(TypedDict):
     category: HarmCategory
     probability: HarmProbability
 
     __doc__ = docstring_utils.strip_oneof(glm.SafetyRating.__doc__)
 
 
-class SafetySetting(TypedDict):
+def convert_rating_to_enum(setting):
+    setting["category"] = HarmCategory(setting["category"])
+    setting["probability"] = HarmProbability(setting["probability"])
+
+
+class SafetySettingDict(TypedDict):
     category: HarmCategory
     threshold: HarmBlockThreshold
 
     __doc__ = docstring_utils.strip_oneof(glm.SafetySetting.__doc__)
 
 
+def convert_setting_to_enum(setting):
+    setting["category"] = HarmCategory(setting["category"])
+    setting["threshold"] = HarmBlockThreshold(setting["threshold"])
+
+
 class SafetyFeedbackDict(TypedDict):
     rating: SafetyRatingDict
-    setting: SafetySetting
+    setting: SafetySettingDict
 
     __doc__ = docstring_utils.strip_oneof(glm.SafetyFeedback.__doc__)
+
+
+def convert_safety_feedback_to_enums(safety_feedback):
+    for sf in safety_feedback:
+        convert_rating_to_enum(sf["rating"])
+        convert_setting_to_enum(sf["setting"])

google/generativeai/types/text_types.py

@@ -32,7 +32,7 @@ class TextCompletion(TypedDict, total=False):
 
 @dataclasses.dataclass(init=False)
 class Completion(abc.ABC):
-    """The result of the `1 given a prompt from the model.
+    """The result returned by `generativeai.generate_text`.
 
     Use `GenerateTextResponse.candidates` to access all the completions generated by the model.
 
@@ -43,9 +43,10 @@ class Completion(abc.ABC):
           Either Unspecified, Safety, or Other. See `types.ContentFilter`.
         safety_feedback: Indicates which safety settings blocked content in this result.
     """
+
     candidates: List[TextCompletion]
     result: Optional[str]
-    filters: Optional[list[safety_types.ContentFilter]]
+    filters: Optional[list[safety_types.ContentFilterDict]]
     safety_feedback: Optional[list[safety_types.SafetyFeedbackDict]]
 
     def to_dict(self) -> Dict[str, Any]:

setup.py

@@ -34,9 +34,7 @@
 else:
     release_status = "Development Status :: 5 - Production/Stable"
 
-dependencies = [
-    "google-ai-generativelanguage==0.2.0"
-]
+dependencies = ["google-ai-generativelanguage==0.2.0"]
 
 extras_require = {
     "dev": [

tests/test_discuss.py

@@ -21,6 +21,8 @@
 from google.generativeai import discuss
 from google.generativeai import client
 import google.generativeai as genai
+from google.generativeai.types import safety_types
+
 from absl.testing import absltest
 from absl.testing import parameterized
 
@@ -36,12 +38,12 @@ def setUp(self):
         self.observed_request = None
 
         self.mock_response = glm.GenerateMessageResponse(
-                candidates=[
-                    glm.Message(content="a", author="1"),
-                    glm.Message(content="b", author="1"),
-                    glm.Message(content="c", author="1"),
-                ],
-            )
+            candidates=[
+                glm.Message(content="a", author="1"),
+                glm.Message(content="b", author="1"),
+                glm.Message(content="c", author="1"),
+            ],
+        )
 
         def fake_generate_message(
             request: glm.GenerateMessageRequest,
@@ -274,32 +276,39 @@ def test_reply(self, kwargs):
         response = response.reply("again")
 
     def test_receive_and_reply_with_filters(self):
-
         self.mock_response = mock_response = glm.GenerateMessageResponse(
             candidates=[glm.Message(content="a", author="1")],
             filters=[
-                glm.ContentFilter(reason=glm.ContentFilter.BlockedReason.SAFETY, message='unsafe'),
-                glm.ContentFilter(reason=glm.ContentFilter.BlockedReason.OTHER),]
+                glm.ContentFilter(
+                    reason=safety_types.BlockedReason.SAFETY, message="unsafe"
+                ),
+                glm.ContentFilter(reason=safety_types.BlockedReason.OTHER),
+            ],
         )
         response = discuss.chat(messages="do filters work?")
 
         filters = response.filters
         self.assertLen(filters, 2)
-        self.assertIsInstance(filters[0]['reason'], glm.ContentFilter.BlockedReason)
-        self.assertEquals(filters[0]['reason'], glm.ContentFilter.BlockedReason.SAFETY)
-        self.assertEquals(filters[0]['message'], 'unsafe')
+        self.assertIsInstance(filters[0]["reason"], safety_types.BlockedReason)
+        self.assertEqual(filters[0]["reason"], safety_types.BlockedReason.SAFETY)
+        self.assertEqual(filters[0]["message"], "unsafe")
 
         self.mock_response = glm.GenerateMessageResponse(
             candidates=[glm.Message(content="a", author="1")],
             filters=[
-                glm.ContentFilter(reason=glm.ContentFilter.BlockedReason.BLOCKED_REASON_UNSPECIFIED)]
+                glm.ContentFilter(
+                    reason=safety_types.BlockedReason.BLOCKED_REASON_UNSPECIFIED
+                )
+            ],
         )
 
-        response = response.reply('Does reply work?')
+        response = response.reply("Does reply work?")
         filters = response.filters
         self.assertLen(filters, 1)
-        self.assertIsInstance(filters[0]['reason'], glm.ContentFilter.BlockedReason)
-        self.assertEquals(filters[0]['reason'], glm.ContentFilter.BlockedReason.BLOCKED_REASON_UNSPECIFIED)
+        self.assertIsInstance(filters[0]["reason"], safety_types.BlockedReason)
+        self.assertEqual(
+            filters[0]["reason"], safety_types.BlockedReason.BLOCKED_REASON_UNSPECIFIED
+        )
 
 
 if __name__ == "__main__":

tests/test_discuss_async.py

@@ -17,8 +17,8 @@
 import unittest
 
 if sys.version_info < (3, 11):
-  import asynctest
-  from asynctest import mock as async_mock
+    import asynctest
+    from asynctest import mock as async_mock
 
 import google.ai.generativelanguage as glm
 
@@ -29,11 +29,16 @@
 bases = (parameterized.TestCase,)
 
 if sys.version_info < (3, 11):
-  bases = bases + (asynctest.TestCase,)
+    bases = bases + (asynctest.TestCase,)
+
+unittest.skipIf(
+    sys.version_info >= (3, 11), "asynctest is not suported on python 3.11+"
+)
+
 
-unittest.skipIf(sys.version_info >= (3,11), "asynctest is not suported on python 3.11+")
 class AsyncTests(*bases):
     if sys.version_info < (3, 11):
+
         async def test_chat_async(self):
             client = async_mock.MagicMock()