Changing safetySettings to map of harm category and harm threshold

langchain4j · Jan 3, 2025 · 816652c · 816652c
1 parent 9c167d1
commit 816652c
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 17 deletions.
diff --git a/...i-spring-boot-starter/src/main/java/dev/langchain4j/googleaigemini/spring/AutoConfig.java b/...i-spring-boot-starter/src/main/java/dev/langchain4j/googleaigemini/spring/AutoConfig.java
@@ -97,16 +97,15 @@ private GeminiMode checkGeminiModeForNull(GeminiFunctionCallingConfig geminiFunc
         return geminiFunctionCallingConfig.geminiMode();
     }
 
-    private Map<GeminiHarmCategory,GeminiHarmBlockThreshold> checkSafetySettingForNull(GeminiSafetySetting safetySetting) {
-        if(safetySetting==null){
+    private Map<GeminiHarmCategory,GeminiHarmBlockThreshold> checkSafetySettingForNull(Map<GeminiHarmCategory,GeminiHarmBlockThreshold> safetySetting) {
+        if(safetySetting==null || safetySetting.isEmpty()){
             Map<GeminiHarmCategory,GeminiHarmBlockThreshold> defaultMap= new HashMap<>();
             defaultMap.put(HARM_CATEGORY_CIVIC_INTEGRITY,HARM_BLOCK_THRESHOLD_UNSPECIFIED);
             return defaultMap;
         }
-        return Map.of(
-                safetySetting.geminiHarmCategory(),
-                safetySetting.geminiHarmBlockThreshold()
-        );
+        Map<GeminiHarmCategory,GeminiHarmBlockThreshold> userMap= new HashMap<>();
+        safetySetting.keySet().forEach(key-> userMap.put(key,safetySetting.get(key)));
+        return userMap;
     }
 
 }
diff --git a/...boot-starter/src/main/java/dev/langchain4j/googleaigemini/spring/ChatModelProperties.java b/...boot-starter/src/main/java/dev/langchain4j/googleaigemini/spring/ChatModelProperties.java
@@ -1,8 +1,11 @@
 package dev.langchain4j.googleaigemini.spring;
 
 import dev.langchain4j.model.chat.request.ResponseFormat;
+import dev.langchain4j.model.googleai.GeminiHarmBlockThreshold;
+import dev.langchain4j.model.googleai.GeminiHarmCategory;
 
 import java.time.Duration;
+import java.util.Map;
 
 public record ChatModelProperties(
         String apiKey,
@@ -15,7 +18,7 @@ public record ChatModelProperties(
         Boolean logRequestsAndResponses,
         Integer maxRetries,
         Duration timeout,
-        GeminiSafetySetting safetySetting,
+        Map<GeminiHarmCategory,GeminiHarmBlockThreshold> safetySetting,
         GeminiFunctionCallingConfig functionCallingConfig
 ) {
 }
diff --git a/...boot-starter/src/main/java/dev/langchain4j/googleaigemini/spring/GeminiSafetySetting.java b/...boot-starter/src/main/java/dev/langchain4j/googleaigemini/spring/GeminiSafetySetting.java