Quick safety filtering: Allow safety_settings="block_none"

google/generativeai/types/safety_types.py

@@ -219,13 +219,25 @@ 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(_NEW_HARM_CATEGORIES.values())
+    return {category: block_threshold for category in set(_NEW_HARM_CATEGORIES.values())}
 
 
 def to_easy_safety_dict(settings: SafetySettingOptions, harm_category_set) -> EasySafetySettingDict:
     if settings is None:
         return {}
-    elif isinstance(settings, Mapping):
+
+    if harm_category_set == "new" and isinstance(settings, (int, str, HarmBlockThreshold)):
+        settings = _expand_block_threshold(settings)
+
+    if isinstance(settings, Mapping):
         return {
             to_harm_category(key, harm_category_set): to_block_threshold(value)
             for key, value in settings.items()
@@ -243,6 +255,10 @@ def normalize_safety_settings(
 ) -> list[SafetySettingDict] | None:
     if settings is None:
         return None
+
+    if harm_category_set == "new" and isinstance(settings, (int, str, HarmBlockThreshold)):
+        settings = _expand_block_threshold(settings)
+
     if isinstance(settings, Mapping):
         return [
             {

tests/test_generative_models.py

@@ -181,7 +181,7 @@ def test_generation_config_overwrite(self, config1, config2):
     )
     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!"),
@@ -198,7 +198,7 @@ def test_safety_overwrite(self, safe1, safe2):
             glm.SafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
         )
 
-        _ = model.generate_content("hello", safety_settings={"danger": "high"})
+        _ = model.generate_content("hello", safety_settings=safe2)
         self.assertEqual(
             self.observed_requests[-1].safety_settings[0].category,
             glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,

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, "new")
+        self.assertEqual(
+            setting[glm.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT],
+            glm.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        )
+
+
+if __name__ == "__main__":
+    absltest.main()