From d0d2dd6bab0166e8f8f063652114c0bc0d8385fb Mon Sep 17 00:00:00 2001
From: Andre Bandarra <andreban@gmail.com>
Date: Tue, 30 Jul 2024 15:04:16 +0100
Subject: [PATCH] Make safety rating scores optional

- The API doesn't always return `probabilityScore` and
`severityScore` for a safety rating, causing the parsing to fail
- This commit makes those fields optional to fix that.
---
 src/types/generate_content.rs | 51 +++++++++++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 2 deletions(-)

diff --git a/src/types/generate_content.rs b/src/types/generate_content.rs
index 4547faa..baac2d9 100644
--- a/src/types/generate_content.rs
+++ b/src/types/generate_content.rs
@@ -88,9 +88,9 @@ pub struct CitationMetadata {
 pub struct SafetyRating {
     pub category: String,
     pub probability: String,
-    pub probability_score: f32,
+    pub probability_score: Option<f32>,
     pub severity: String,
-    pub severity_score: f32,
+    pub severity_score: Option<f32>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -296,4 +296,51 @@ mod tests {
       }"#;
         serde_json::from_str::<GenerateContentResponse>(input).unwrap();
     }
+
+    #[test]
+    fn parses_safety_rating_without_scores() {
+        let input = r#"{
+          "candidates": [
+            {
+              "content": {
+                "role": "model",
+                "parts": [
+                  {
+                    "text": "Return text"
+                  }
+                ]
+              },
+              "finishReason": "STOP",
+              "safetyRatings": [
+                {
+                  "category": "HARM_CATEGORY_HATE_SPEECH",
+                  "probability": "NEGLIGIBLE",
+                  "severity": "HARM_SEVERITY_NEGLIGIBLE"
+                },
+                {
+                  "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                  "probability": "NEGLIGIBLE",
+                  "severity": "HARM_SEVERITY_NEGLIGIBLE"
+                },
+                {
+                  "category": "HARM_CATEGORY_HARASSMENT",
+                  "probability": "NEGLIGIBLE",
+                  "severity": "HARM_SEVERITY_NEGLIGIBLE"
+                },
+                {
+                  "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                  "probability": "NEGLIGIBLE",
+                  "severity": "HARM_SEVERITY_NEGLIGIBLE"
+                }
+              ]
+            }
+          ],
+          "usageMetadata": {
+            "promptTokenCount": 5492,
+            "candidatesTokenCount": 1256,
+            "totalTokenCount": 6748
+          }
+        }"#;
+        serde_json::from_str::<GenerateContentResponse>(input).unwrap();
+    }
 }