Changing the safteySetting and toolConfig and writing test

langchain4j · Nov 30, 2024 · 7554b25 · 7554b25
1 parent 4a27ef1
commit 7554b25
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
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
@@ -54,6 +54,12 @@ StreamingChatLanguageModel googleAiGeminiStreamingChatModel(Properties propertie
                 .topK(chatModelProperties.getTopK())
                 .responseFormat(chatModelProperties.getResponseFormat())
                 .logRequestsAndResponses(chatModelProperties.getLogRequestsAndResponses())
+                .safetySettings(
+                        Map.of(chatModelProperties.getSafetySetting().getGeminiHarmCategory(),
+                                chatModelProperties.getSafetySetting().getGeminiHarmBlockThreshold()))
+                .toolConfig(
+                        chatModelProperties.getFunctionCallingConfig().getGeminiMode(),
+                        chatModelProperties.getFunctionCallingConfig().getAllowedFunctionNames().toArray(new String[0]))
                 .build();
     }
 

diff --git a/...spring-boot-starter/src/test/java/dev/langchain4j/googleaigemini/spring/AutoConfigIT.java b/...spring-boot-starter/src/test/java/dev/langchain4j/googleaigemini/spring/AutoConfigIT.java
@@ -62,7 +62,11 @@ void provide_streaming_chat_model() {
                         "langchain4j.google-ai-gemini.streamingChatModel.temperature=0.7",
                         "langchain4j.google-ai-gemini.streamingChatModel.topP=0.9",
                         "langchain4j.google-ai-gemini.streamingChatModel.topK=40",
-                        "langchain4j.google-ai-gemini.streamingChatModel.maxOutputTokens=100"
+                        "langchain4j.google-ai-gemini.streamingChatModel.maxOutputTokens=100",
+                        "langchain4j.google-ai-gemini.chatModel.safetySetting.gemini-harm-category=HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                        "langchain4j.google-ai-gemini.chatModel.safetySetting.gemini-harm-block-threshold=HARM_BLOCK_THRESHOLD_UNSPECIFIED",
+                        "langchain4j.google-ai-gemini.chatModel.functionCallingConfig.gemini-mode=ANY",
+                        "langchain4j.google-ai-gemini.chatModel.functionCallingConfig.allowed-function-names=allowCodeExecution,includeCodeExecutionOutput"
                 )
                 .run(context -> {
                     StreamingChatLanguageModel streamingChatLanguageModel = context.getBean(StreamingChatLanguageModel.class);