Test candidates['safety_ratings'].

Author: MarkDaoust

google/generativeai/client.py

@@ -45,7 +45,7 @@ def configure(
     # but that seems rare. Users that need it can just switch to the low level API.
     transport: Union[str, None] = None,
     client_options: Union[client_options_lib.ClientOptions, dict, None] = None,
-    client_info: Optional[gapic_v1.client_info.ClientInfo] = None
+    client_info: Optional[gapic_v1.client_info.ClientInfo] = None,
 ):
     """Captures default client configuration.
 
@@ -86,13 +86,13 @@ def configure(
 
     user_agent = f"{USER_AGENT}/{version.__version__}"
     if client_info:
-      # Be respectful of any existing agent setting.
-      if client_info.user_agent:
-        client_info.user_agent += f" {user_agent}"
-      else:
-        client_info.user_agent = user_agent
+        # Be respectful of any existing agent setting.
+        if client_info.user_agent:
+            client_info.user_agent += f" {user_agent}"
+        else:
+            client_info.user_agent = user_agent
     else:
-      client_info = gapic_v1.client_info.ClientInfo(user_agent=user_agent)
+        client_info = gapic_v1.client_info.ClientInfo(user_agent=user_agent)
 
     new_default_client_config = {
         "credentials": credentials,

google/generativeai/discuss.py

@@ -455,7 +455,7 @@ def _build_chat_response(
     response = type(response).to_dict(response)
     response.pop("messages")
 
-    safety_types.convert_filters_to_enums(response["filters"])
+    response["filters"] = safety_types.convert_filters_to_enums(response["filters"])
 
     if response["candidates"]:
         last = response["candidates"][0]

google/generativeai/text.py

@@ -159,8 +159,11 @@ 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"])
+    response["filters"] = safety_types.convert_filters_to_enums(response["filters"])
+    response["safety_feedback"] = safety_types.convert_safety_feedback_to_enums(
+        response["safety_feedback"]
+    )
+    response['candidates'] = safety_types.convert_candidate_enums(response['candidates'])
 
     return Completion(_client=client, **response)
 

google/generativeai/types/safety_types.py

@@ -16,7 +16,7 @@
 import enum
 from google.ai import generativelanguage as glm
 from google.generativeai import docstring_utils
-from typing import TypedDict
+from typing import Iterable, List, TypedDict
 
 __all__ = [
     "HarmCategory",
@@ -43,9 +43,13 @@ class ContentFilterDict(TypedDict):
     __doc__ = docstring_utils.strip_oneof(glm.ContentFilter.__doc__)
 
 
-def convert_filters_to_enums(filters):
+def convert_filters_to_enums(filters: Iterable[dict]) -> List[ContentFilterDict]:
+    result = []
     for f in filters:
+        f = f.copy()
         f["reason"] = BlockedReason(f["reason"])
+        result.append(f)
+    return result
 
 
 class SafetyRatingDict(TypedDict):
@@ -55,9 +59,18 @@ class SafetyRatingDict(TypedDict):
     __doc__ = docstring_utils.strip_oneof(glm.SafetyRating.__doc__)
 
 
-def convert_rating_to_enum(setting):
-    setting["category"] = HarmCategory(setting["category"])
-    setting["probability"] = HarmProbability(setting["probability"])
+def convert_rating_to_enum(rating: dict) -> SafetyRatingDict:
+    return {
+        "category": HarmCategory(rating["category"]),
+        "probability": HarmProbability(rating["probability"]),
+    }
+
+
+def convert_ratings_to_enum(ratings: Iterable[dict]) -> List[SafetyRatingDict]:
+    result = []
+    for r in ratings:
+        result.append(convert_rating_to_enum(r))
+    return result
 
 
 class SafetySettingDict(TypedDict):
@@ -67,9 +80,11 @@ class SafetySettingDict(TypedDict):
     __doc__ = docstring_utils.strip_oneof(glm.SafetySetting.__doc__)
 
 
-def convert_setting_to_enum(setting):
-    setting["category"] = HarmCategory(setting["category"])
-    setting["threshold"] = HarmBlockThreshold(setting["threshold"])
+def convert_setting_to_enum(setting: dict) -> SafetySettingDict:
+    return {
+        "category": HarmCategory(setting["category"]),
+        "threshold": HarmBlockThreshold(setting["threshold"]),
+    }
 
 
 class SafetyFeedbackDict(TypedDict):
@@ -79,7 +94,24 @@ class SafetyFeedbackDict(TypedDict):
     __doc__ = docstring_utils.strip_oneof(glm.SafetyFeedback.__doc__)
 
 
-def convert_safety_feedback_to_enums(safety_feedback):
+def convert_safety_feedback_to_enums(
+    safety_feedback: Iterable[dict],
+) -> List[SafetyFeedbackDict]:
+    result = []
     for sf in safety_feedback:
-        convert_rating_to_enum(sf["rating"])
-        convert_setting_to_enum(sf["setting"])
+        result.append(
+            {
+                "rating": convert_rating_to_enum(sf["rating"]),
+                "setting": convert_setting_to_enum(sf["setting"]),
+            }
+        )
+    return result
+
+
+def convert_candidate_enums(candidates):
+    result = []
+    for candidate in candidates:
+        candidate = candidate.copy()
+        candidate['safety_ratings'] = convert_ratings_to_enum(candidate['safety_ratings'])
+        result.append(candidate)
+    return result

tests/test_text.py

@@ -223,9 +223,33 @@ def test_safety_feedback(self):
             safety_types.HarmCategory.HARM_CATEGORY_MEDICAL,
         )
 
-    #def test_candidate_safety_feedback(self):
+    def test_candidate_safety_feedback(self):
+        self.mock_response = glm.GenerateTextResponse(
+            candidates=[
+                {
+                    "output": "hello",
+                    "safety_ratings": [
+                        {
+                            "category": safety_types.HarmCategory.HARM_CATEGORY_MEDICAL,
+                            "probability": safety_types.HarmProbability.HIGH,
+                        },
+                        {
+                            "category": safety_types.HarmCategory.HARM_CATEGORY_VIOLENCE,
+                            "probability": safety_types.HarmProbability.LOW,
+                        },
+                    ],
+                }
+            ]
+        )
+
+        result = text_service.generate_text(prompt="Write a story from the ER.")
+        self.assertIsInstance(result.candidates[0]['safety_ratings'][0]['category'], safety_types.HarmCategory)
+        self.assertEqual(result.candidates[0]['safety_ratings'][0]['category'], safety_types.HarmCategory.HARM_CATEGORY_MEDICAL)
+
+        self.assertIsInstance(result.candidates[0]['safety_ratings'][0]['probability'], safety_types.HarmProbability)
+        self.assertEqual(result.candidates[0]['safety_ratings'][0]['probability'], safety_types.HarmProbability.HIGH)
 
-    #def test_candidate_citations(self):
+    # def test_candidate_citations(self):
 
 
 if __name__ == "__main__":