Test text candidate citations and safety ratings.

Author: MarkDaoust

google/generativeai/text.py

@@ -163,7 +163,9 @@ def _generate_response(
     response["safety_feedback"] = safety_types.convert_safety_feedback_to_enums(
         response["safety_feedback"]
     )
-    response['candidates'] = safety_types.convert_candidate_enums(response['candidates'])
+    response["candidates"] = safety_types.convert_candidate_enums(
+        response["candidates"]
+    )
 
     return Completion(_client=client, **response)
 

google/generativeai/types/safety_types.py

@@ -112,6 +112,8 @@ def convert_candidate_enums(candidates):
     result = []
     for candidate in candidates:
         candidate = candidate.copy()
-        candidate['safety_ratings'] = convert_ratings_to_enum(candidate['safety_ratings'])
+        candidate["safety_ratings"] = convert_ratings_to_enum(
+            candidate["safety_ratings"]
+        )
         result.append(candidate)
-    return result
+    return result

tests/test_text.py

@@ -167,7 +167,6 @@ def test_safety_settings(self):
             ],
         )
 
-        # Just make sure it made it into the request object.
         self.assertEqual(
             self.observed_request.safety_settings[0].category,
             safety_types.HarmCategory.HARM_CATEGORY_MEDICAL,
@@ -243,13 +242,48 @@ def test_candidate_safety_feedback(self):
         )
 
         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]["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)
+        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):
+        self.mock_response = glm.GenerateTextResponse(
+            candidates=[
+                {
+                    "output": "Hello Google!",
+                    "citation_metadata": {
+                        "citation_sources": [
+                            {
+                                "start_index": 6,
+                                "end_index": 12,
+                                "uri": "https://google.com",
+                            }
+                        ]
+                    },
+                }
+            ]
+        )
+        result = text_service.generate_text(prompt="Hi my name is Google")
+        self.assertEqual(
+            result.candidates[0]["citation_metadata"]["citation_sources"][0][
+                "start_index"
+            ],
+            6,
+        )
 
 
 if __name__ == "__main__":