feat: GenAI - Added support for SafetySetting.method (probability or severity)

Ark-kun · copybara-github · commit 317ab8f7499b · 2024-04-02T16:39:41.000-07:00
PiperOrigin-RevId: 621331507
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
@@ -307,6 +307,18 @@ def test_generate_content(self, generative_models: generative_models):
                 max_output_tokens=200,
                 stop_sequences=["\n\n\n"],
             ),
+            safety_settings=[
+                generative_models.SafetySetting(
+                    category=generative_models.SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                    threshold=generative_models.SafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+                    method=generative_models.SafetySetting.HarmBlockMethod.SEVERITY,
+                ),
+                generative_models.SafetySetting(
+                    category=generative_models.SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+                    threshold=generative_models.SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,
+                    method=generative_models.SafetySetting.HarmBlockMethod.PROBABILITY,
+                ),
+            ],
         )
         assert response2.text
 
diff --git a/vertexai/generative_models/__init__.py b/vertexai/generative_models/__init__.py
@@ -30,6 +30,7 @@
     Image,
     Part,
     ResponseValidationError,
+    SafetySetting,
     Tool,
 )
 
@@ -47,5 +48,6 @@
     "Image",
     "Part",
     "ResponseValidationError",
+    "SafetySetting",
     "Tool",
 ]
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -86,7 +86,7 @@
 ]
 
 SafetySettingsType = Union[
-    List[gapic_content_types.SafetySetting],
+    List["SafetySetting"],
     Dict[
         gapic_content_types.HarmCategory,
         gapic_content_types.SafetySetting.HarmBlockThreshold,
@@ -258,17 +258,20 @@ def _prepare_request(
                 raise TypeError(
                     "generation_config must either be a GenerationConfig object or a dictionary representation of it."
                 )
+
         gapic_safety_settings = None
         if safety_settings:
             if isinstance(safety_settings, Sequence):
-                if not all(
-                    isinstance(safety_setting, gapic_content_types.SafetySetting)
-                    for safety_setting in safety_settings
-                ):
-                    raise TypeError(
-                        "When passing a list with SafetySettings objects, every item in a list must be a SafetySetting object."
-                    )
-                gapic_safety_settings = safety_settings
+                gapic_safety_settings = []
+                for safety_setting in safety_settings:
+                    if isinstance(safety_setting, gapic_content_types.SafetySetting):
+                        gapic_safety_settings.append(safety_setting)
+                    elif isinstance(safety_setting, SafetySetting):
+                        gapic_safety_settings.append(safety_setting._raw_safety_setting)
+                    else:
+                        raise TypeError(
+                            "When passing a list with SafetySettings objects, every item in a list must be a SafetySetting object."
+                        )
             elif isinstance(safety_settings, dict):
                 gapic_safety_settings = [
                     gapic_content_types.SafetySetting(
@@ -283,6 +286,7 @@ def _prepare_request(
                 raise TypeError(
                     "safety_settings must either be a list of SafetySettings objects or a dictionary mapping from HarmCategory to HarmBlockThreshold."
                 )
+
         gapic_tools = None
         if tools:
             gapic_tools = []
@@ -1738,6 +1742,61 @@ def _image(self) -> "Image":
         return Image.from_bytes(data=self._raw_part.inline_data.data)
 
 
+class SafetySetting:
+    """Parameters for the generation."""
+
+    HarmCategory = gapic_content_types.HarmCategory
+    HarmBlockMethod = gapic_content_types.SafetySetting.HarmBlockMethod
+    HarmBlockThreshold = gapic_content_types.SafetySetting.HarmBlockThreshold
+
+    def __init__(
+        self,
+        *,
+        category: "SafetySetting.HarmCategory",
+        threshold: "SafetySetting.HarmBlockThreshold",
+        method: Optional["SafetySetting.HarmBlockMethod"] = None,
+    ):
+        r"""Safety settings.
+
+        Args:
+            category: Harm category.
+            threshold: The harm block threshold.
+            method: Specify if the threshold is used for probability or severity
+                score. If not specified, the threshold is used for probability
+                score.
+        """
+        self._raw_safety_setting = gapic_content_types.SafetySetting(
+            category=category,
+            threshold=threshold,
+            method=method,
+        )
+
+    @classmethod
+    def _from_gapic(
+        cls,
+        raw_safety_setting: gapic_content_types.SafetySetting,
+    ) -> "SafetySetting":
+        response = cls(
+            category=raw_safety_setting.category,
+            threshold=raw_safety_setting.threshold,
+        )
+        response._raw_safety_setting = raw_safety_setting
+        return response
+
+    @classmethod
+    def from_dict(cls, safety_setting_dict: Dict[str, Any]) -> "SafetySetting":
+        raw_safety_setting = gapic_content_types.SafetySetting(
+            safety_setting_dict
+        )
+        return cls._from_gapic(raw_safety_setting=raw_safety_setting)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return type(self._raw_safety_setting).to_dict(self._raw_safety_setting)
+
+    def __repr__(self):
+        return self._raw_safety_setting.__repr__()
+
+
 class grounding:  # pylint: disable=invalid-name
     """Grounding namespace."""
 
diff --git a/vertexai/preview/generative_models.py b/vertexai/preview/generative_models.py
@@ -34,6 +34,7 @@
     Part,
     ResponseBlockedError,
     ResponseValidationError,
+    SafetySetting,
     Tool,
 )
 
@@ -64,6 +65,7 @@ class ChatSession(_PreviewChatSession):
     "Part",
     "ResponseBlockedError",
     "ResponseValidationError",
+    "SafetySetting",
     "Tool",
     #
 ]

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@`
`30`	`30`	`Image,`
`31`	`31`	`Part,`
`32`	`32`	`ResponseValidationError,`
	`33`	`+ SafetySetting,`
`33`	`34`	`Tool,`
`34`	`35`	`)`
`35`	`36`
`@@ -47,5 +48,6 @@`
`47`	`48`	`"Image",`
`48`	`49`	`"Part",`
`49`	`50`	`"ResponseValidationError",`
	`51`	`+ "SafetySetting",`
`50`	`52`	`"Tool",`
`51`	`53`	`]`
Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,7 @@`
`34`	`34`	`Part,`
`35`	`35`	`ResponseBlockedError,`
`36`	`36`	`ResponseValidationError,`
	`37`	`+ SafetySetting,`
`37`	`38`	`Tool,`
`38`	`39`	`)`
`39`	`40`
`@@ -64,6 +65,7 @@ class ChatSession(_PreviewChatSession):`
`64`	`65`	`"Part",`
`65`	`66`	`"ResponseBlockedError",`
`66`	`67`	`"ResponseValidationError",`
	`68`	`+ "SafetySetting",`
`67`	`69`	`"Tool",`
`68`	`70`	`#`
`69`	`71`	`]`