From 9b283024b47f57ad7bbf032fa015eab42f8162a9 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 6 May 2024 08:19:42 +0800
Subject: feat: extract vertexai-claude client (#485)

---
 src/client/gemini.rs | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

(limited to 'src/client/gemini.rs')

diff --git a/src/client/gemini.rs b/src/client/gemini.rs
index 8f6a76d..28fc93e 100644
--- a/src/client/gemini.rs
+++ b/src/client/gemini.rs
@@ -11,7 +11,8 @@ const API_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/models/
 pub struct GeminiConfig {
     pub name: Option<String>,
     pub api_key: Option<String>,
-    pub block_threshold: Option<String>,
+    #[serde(rename = "safetySettings")]
+    pub safety_settings: Option<serde_json::Value>,
     #[serde(default)]
     pub models: Vec<ModelConfig>,
     pub extra: Option<ExtraConfig>,
@@ -31,9 +32,7 @@ impl GeminiClient {
             false => "generateContent",
         };
 
-        let block_threshold = self.config.block_threshold.clone();
-
-        let body = gemini_build_body(data, &self.model, block_threshold)?;
+        let body = gemini_build_body(data, &self.model, self.config.safety_settings.clone())?;
 
         let model = &self.model.name;
 
-- 
cgit v1.2.3