allow missing safetyRatings

jelni · Dec 27, 2023 · bdd0e9e · bdd0e9e
1 parent 6e68537
commit bdd0e9e
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 71 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/apis/makersuite.rs b/src/apis/makersuite.rs
@@ -49,6 +49,7 @@ struct GenerationConfig {
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct GenerateContentResponse {
+    #[serde(default)]
     pub candidates: Vec<Candidate>,
     pub prompt_feedback: Option<PromptFeedback>,
 }
@@ -96,13 +97,14 @@ pub struct CitationSource {
 #[serde(rename_all = "camelCase")]
 pub struct PromptFeedback {
     pub block_reason: Option<String>,
-    pub safety_ratings: Vec<SafetyRating>,
+    pub safety_ratings: Option<Vec<SafetyRating>>,
 }
 
 #[derive(Debug, Deserialize)]
 pub struct SafetyRating {
     pub category: String,
-    pub blocked: Option<bool>,
+    #[serde(default)]
+    pub blocked: bool,
 }
 
 #[derive(Deserialize)]

diff --git a/src/commands/makersuite.rs b/src/commands/makersuite.rs
@@ -92,23 +92,18 @@ impl CommandTrait for GoogleGemini {
         };
 
         if let Some(prompt_feedback) = response.prompt_feedback {
-            if let Some(block_reason) = prompt_feedback.block_reason {
-                if block_reason == "SAFETY" {
-                    let reasons = prompt_feedback
-                        .safety_ratings
-                        .into_iter()
-                        .filter(|safety_rating| safety_rating.blocked.unwrap_or_default())
-                        .map(|safety_rating| safety_rating.category)
-                        .collect::<Vec<_>>()
-                        .join(", ");
-
-                    return Err(CommandError::Custom(format!(
-                        "request blocked by Google: {reasons}."
-                    )));
-                };
-
-                return Err(CommandError::Custom("request blocked by Google.".into()));
-            };
+            if let Some(safety_ratings) = prompt_feedback.safety_ratings {
+                let reasons = safety_ratings
+                    .into_iter()
+                    .filter(|safety_rating| safety_rating.blocked)
+                    .map(|safety_rating| safety_rating.category)
+                    .collect::<Vec<_>>()
+                    .join(", ");
+
+                return Err(CommandError::Custom(format!("request blocked by Google: {reasons}.")));
+            }
+
+            return Err(CommandError::Custom("request blocked by Google.".into()));
         }
 
         let Some(candidate) = response.candidates.into_iter().next() else {