Quick safety filtering: Allow safety_settings="block_none" (#347)

* allow safety_settings='off'

Change-Id: Ica10b399177301073424a98cb3a8b0736dc216b4

* Fix tests.

Change-Id: I06cfd07397e984b9fb757b2831b419eefb8aff98

* license

Change-Id: Ifa4843831b9c1479198c2b45c5b5abad8410f448

* format

Change-Id: I534837c309121cda9c8947acdd6c126c9c730d62

* add test

Change-Id: I9bce66322d64b3d6296d4db7cc0a7b7b9a78763b

Author: MarkDaoust

google/generativeai/types/safety_types.py

@@ -201,25 +201,52 @@ class LooseSafetySettingDict(TypedDict):
 EasySafetySetting = Mapping[HarmCategoryOptions, HarmBlockThresholdOptions]
 EasySafetySettingDict = dict[HarmCategoryOptions, HarmBlockThresholdOptions]
 
-SafetySettingOptions = Union[EasySafetySetting, Iterable[LooseSafetySettingDict], None]
+SafetySettingOptions = Union[
+    HarmBlockThresholdOptions, EasySafetySetting, Iterable[LooseSafetySettingDict], None
+]
+
+
+def _expand_block_threshold(block_threshold: HarmBlockThresholdOptions):
+    block_threshold = to_block_threshold(block_threshold)
+    set(_HARM_CATEGORIES.values())
+    return {category: block_threshold for category in set(_HARM_CATEGORIES.values())}
 
 
 def to_easy_safety_dict(settings: SafetySettingOptions) -> EasySafetySettingDict:
     if settings is None:
         return {}
-    elif isinstance(settings, Mapping):
+
+    if isinstance(settings, (int, str, HarmBlockThreshold)):
+        settings = _expand_block_threshold(settings)
+
+    if isinstance(settings, Mapping):
         return {to_harm_category(key): to_block_threshold(value) for key, value in settings.items()}
+
     else:  # Iterable
-        return {
-            to_harm_category(d["category"]): to_block_threshold(d["threshold"]) for d in settings
-        }
+        result = {}
+        for setting in settings:
+            if isinstance(setting, glm.SafetySetting):
+                result[to_harm_category(setting.category)] = to_block_threshold(setting.threshold)
+            elif isinstance(setting, dict):
+                result[to_harm_category(setting["category"])] = to_block_threshold(
+                    setting["threshold"]
+                )
+            else:
+                raise ValueError(
+                    f"Could not understand safety setting:\n  {type(setting)=}\n  {setting=}"
+                )
+        return result
 
 
 def normalize_safety_settings(
     settings: SafetySettingOptions,
 ) -> list[SafetySettingDict] | None:
     if settings is None:
         return None
+
+    if isinstance(settings, (int, str, HarmBlockThreshold)):
+        settings = _expand_block_threshold(settings)
+
     if isinstance(settings, Mapping):
         return [
             {

tests/test_generative_models.py

@@ -155,11 +155,12 @@ def test_generation_config_overwrite(self, config1, config2):
 
     @parameterized.named_parameters(
         ["dict", {"danger": "low"}, {"danger": "high"}],
+        ["quick", "low", "high"],
         [
             "list-dict",
             [
                 dict(
-                    category=glm.HarmCategory.HARM_CATEGORY_DANGEROUS,
+                    category=glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
                     threshold=glm.SafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
                 ),
             ],
@@ -171,44 +172,47 @@ def test_generation_config_overwrite(self, config1, config2):
             "object",
             [
                 glm.SafetySetting(
-                    category=glm.HarmCategory.HARM_CATEGORY_DANGEROUS,
+                    category=glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
                     threshold=glm.SafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
                 ),
             ],
             [
                 glm.SafetySetting(
-                    category=glm.HarmCategory.HARM_CATEGORY_DANGEROUS,
-                    threshold=glm.SafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+                    category=glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                    threshold=glm.SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,
                 ),
             ],
         ],
     )
     def test_safety_overwrite(self, safe1, safe2):
         # Safety
-        model = generative_models.GenerativeModel("gemini-pro", safety_settings={"danger": "low"})
+        model = generative_models.GenerativeModel("gemini-pro", safety_settings=safe1)
 
         self.responses["generate_content"] = [
             simple_response(" world!"),
             simple_response(" world!"),
         ]
 
         _ = model.generate_content("hello")
+
+        danger = [
+            s
+            for s in self.observed_requests[-1].safety_settings
+            if s.category == glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT
+        ]
         self.assertEqual(
-            self.observed_requests[-1].safety_settings[0].category,
-            glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
-        )
-        self.assertEqual(
-            self.observed_requests[-1].safety_settings[0].threshold,
+            danger[0].threshold,
             glm.SafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
         )
 
-        _ = model.generate_content("hello", safety_settings={"danger": "high"})
-        self.assertEqual(
-            self.observed_requests[-1].safety_settings[0].category,
-            glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
-        )
+        _ = model.generate_content("hello", safety_settings=safe2)
+        danger = [
+            s
+            for s in self.observed_requests[-1].safety_settings
+            if s.category == glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT
+        ]
         self.assertEqual(
-            self.observed_requests[-1].safety_settings[0].threshold,
+            danger[0].threshold,
             glm.SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,
         )
 

tests/test_safety.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import google.ai.generativelanguage as glm
+from google.generativeai.types import safety_types
+
+
+class SafetyTests(parameterized.TestCase):
+    """Tests are in order with the design doc."""
+
+    @parameterized.named_parameters(
+        ["block_threshold", glm.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE],
+        ["block_threshold2", "medium"],
+        ["block_threshold3", 2],
+        ["dict", {"danger": "medium"}],
+        ["dict2", {"danger": 2}],
+        ["dict3", {"danger": glm.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE}],
+        [
+            "list-dict",
+            [
+                dict(
+                    category=glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                    threshold=glm.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+                ),
+            ],
+        ],
+        [
+            "list-dict2",
+            [
+                dict(category="danger", threshold="med"),
+            ],
+        ],
+    )
+    def test_safety_overwrite(self, setting):
+        setting = safety_types.to_easy_safety_dict(setting)
+        self.assertEqual(
+            setting[glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT],
+            glm.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        )
+
+
+if __name__ == "__main__":
+    absltest.main()