Test filters in chat

Author: MarkDaoust

google/generativeai/discuss.py

@@ -438,6 +438,9 @@ 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,
@@ -453,6 +456,8 @@ def _build_chat_response(
     response = type(response).to_dict(response)
     response.pop("messages")
 
+    _convert_filters_to_enums(response['filters'])
+
     if response["candidates"]:
         last = response["candidates"][0]
     else:

google/generativeai/types/safety_types.py

@@ -22,7 +22,7 @@
     "HarmCategory",
     "HarmProbability",
     "HarmBlockThreshold",
-    "BlockReason",
+    "BlockedReason",
     "ContentFilter",
     "SafetyRatingDict",
     "SafetySetting",
@@ -33,11 +33,11 @@
 HarmCategory = glm.HarmCategory
 HarmProbability = glm.SafetyRating.HarmProbability
 HarmBlockThreshold = glm.SafetySetting.HarmBlockThreshold
-BlockReason = glm.ContentFilter.BlockedReason
+BlockedReason = glm.ContentFilter.BlockedReason
 
 
 class ContentFilter(TypedDict):
-    reason: BlockReason
+    reason: BlockedReason
     message: str
 
     __doc__ = docstring_utils.strip_oneof(glm.ContentFilter.__doc__)

tests/test_discuss.py

@@ -12,6 +12,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import copy
+
 import unittest.mock
 
 import google.ai.generativelanguage as glm
@@ -33,19 +35,22 @@ def setUp(self):
 
         self.observed_request = None
 
-        def fake_generate_message(
-            request: glm.GenerateMessageRequest,
-        ) -> glm.GenerateMessageResponse:
-            self.observed_request = request
-            return glm.GenerateMessageResponse(
-                messages=request.prompt.messages,
+        self.mock_response = glm.GenerateMessageResponse(
                 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,
+        ) -> glm.GenerateMessageResponse:
+            self.observed_request = request
+            response = copy.copy(self.mock_response)
+            response.messages = request.prompt.messages
+            return response
+
         self.client.generate_message = fake_generate_message
 
     @parameterized.named_parameters(
@@ -268,6 +273,34 @@ 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),]
+        )
+        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.mock_response = glm.GenerateMessageResponse(
+            candidates=[glm.Message(content="a", author="1")],
+            filters=[
+                glm.ContentFilter(reason=glm.ContentFilter.BlockedReason.BLOCKED_REASON_UNSPECIFIED)]
+        )
+
+        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)
+
 
 if __name__ == "__main__":
     absltest.main()

tests/test_discuss_async.py

@@ -64,28 +64,28 @@ async def fake_generate_message(
                 client=client,
             )
 
-        self.assertEqual(
-            observed_request,
-            glm.GenerateMessageRequest(
-                model="models/bard",
-                prompt=glm.MessagePrompt(
-                    context="Example Prompt",
-                    examples=[
-                        glm.Example(
-                            input=glm.Message(content="Example from human"),
-                            output=glm.Message(content="Example response from AI"),
-                        )
-                    ],
-                    messages=[glm.Message(author="0", content="Tell me a joke")],
+            self.assertEqual(
+                observed_request,
+                glm.GenerateMessageRequest(
+                    model="models/bard",
+                    prompt=glm.MessagePrompt(
+                        context="Example Prompt",
+                        examples=[
+                            glm.Example(
+                                input=glm.Message(content="Example from human"),
+                                output=glm.Message(content="Example response from AI"),
+                            )
+                        ],
+                        messages=[glm.Message(author="0", content="Tell me a joke")],
+                    ),
+                    temperature=0.75,
+                    candidate_count=1,
                 ),
-                temperature=0.75,
-                candidate_count=1,
-            ),
-        )
-        self.assertEqual(
-            observed_response.candidates,
-            [{"author": "1", "content": "Why did the chicken cross the road?"}],
-        )
+            )
+            self.assertEqual(
+                observed_response.candidates,
+                [{"author": "1", "content": "Why did the chicken cross the road?"}],
+            )
 
 
 if __name__ == "__main__":