From 1d467138eda12fe610ab80fa7bfe71fa79faf178 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Cipriani=20Bandarra?= <andreban@gmail.com>
Date: Sun, 8 Jun 2025 09:24:34 +0100
Subject: [PATCH] Move predict image safety settings to enum

---
 examples/generate_image.rs |  2 ++
 src/types/predict_image.rs | 11 ++++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/examples/generate_image.rs b/examples/generate_image.rs
index 221c40e..f894f95 100644
--- a/examples/generate_image.rs
+++ b/examples/generate_image.rs
@@ -3,6 +3,7 @@ use std::{error::Error, io::Cursor};
 use gemini_rs::prelude::{
     GeminiClient, PersonGeneration, PredictImageRequest, PredictImageRequestParameters,
     PredictImageRequestParametersOutputOptions, PredictImageRequestPrompt,
+    PredictImageSafetySetting,
 };
 use image::{ImageFormat, ImageReader};
 
@@ -35,6 +36,7 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
                 compression_quality: Some(75),
             }),
             person_generation: Some(PersonGeneration::AllowAll),
+            safety_setting: Some(PredictImageSafetySetting::BlockOnlyHigh),
             ..Default::default()
         },
     };
diff --git a/src/types/predict_image.rs b/src/types/predict_image.rs
index f0ebe20..3531a38 100644
--- a/src/types/predict_image.rs
+++ b/src/types/predict_image.rs
@@ -126,7 +126,7 @@ pub struct PredictImageRequestParameters {
     /// Supported by the models `imagen-3.0-generate-001`, `imagen-3.0-fast-generate-001`, and
     /// `imagegeneration@006` only.
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub safety_setting: Option<String>,
+    pub safety_setting: Option<PredictImageSafetySetting>,
 
     /// Add an invisible watermark to the generated images. The default value is `false` for the
     /// `imagegeneration@002` and `imagegeneration@005` models, and `true` for the
@@ -176,3 +176,12 @@ pub enum PersonGeneration {
     AllowAdult,
     AllowAll,
 }
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum PredictImageSafetySetting {
+    BlockLowAndAbove,
+    BlockMediumAndAbove,
+    BlockOnlyHigh,
+    BlockNone,
+}