feat(gemini) Support Safety Settings for VertexAiGeminiChatModel

Author: terussell85

models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java

@@ -453,7 +453,8 @@ GeminiRequest createGeminiRequest(Prompt prompt, VertexAiGeminiChatOptions updat
 		GenerationConfig generationConfig = this.generationConfig;
 
 		var generativeModelBuilder = new GenerativeModel.Builder().setModelName(this.defaultOptions.getModel())
-			.setVertexAi(this.vertexAI);
+			.setVertexAi(this.vertexAI)
+			.setSafetySettings(this.defaultOptions.getSafetySettings());
 
 		if (prompt.getOptions() != null) {
 			if (prompt.getOptions() instanceof FunctionCallingOptions functionCallingOptions) {
@@ -499,6 +500,10 @@ GeminiRequest createGeminiRequest(Prompt prompt, VertexAiGeminiChatOptions updat
 			generativeModelBuilder.setTools(tools);
 		}
 
+		if (prompt.getOptions() instanceof VertexAiGeminiChatOptions options && options.getSafetySettings() != null) {
+			generativeModelBuilder.setSafetySettings(options.getSafetySettings());
+		}
+
 		generativeModelBuilder.setGenerationConfig(generationConfig);
 
 		GenerativeModel generativeModel = generativeModelBuilder.build();

models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatOptions.java

@@ -28,6 +28,7 @@
 import com.fasterxml.jackson.annotation.JsonInclude.Include;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import com.google.cloud.vertexai.api.SafetySetting;
 import org.springframework.ai.model.function.FunctionCallback;
 import org.springframework.ai.model.function.FunctionCallingOptions;
 import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatModel.ChatModel;
@@ -117,6 +118,9 @@ public class VertexAiGeminiChatOptions implements FunctionCallingOptions {
 	@JsonIgnore
 	private boolean googleSearchRetrieval = false;
 
+	@JsonIgnore
+	private List<SafetySetting> safetySettings = new ArrayList<>();
+
 	@JsonIgnore
 	private Boolean proxyToolCalls;
 
@@ -143,6 +147,7 @@ public static VertexAiGeminiChatOptions fromOptions(VertexAiGeminiChatOptions fr
 		options.setFunctions(fromOptions.getFunctions());
 		options.setResponseMimeType(fromOptions.getResponseMimeType());
 		options.setGoogleSearchRetrieval(fromOptions.getGoogleSearchRetrieval());
+		options.setSafetySettings(fromOptions.getSafetySettings());
 		options.setProxyToolCalls(fromOptions.getProxyToolCalls());
 		options.setToolContext(fromOptions.getToolContext());
 		return options;
@@ -269,6 +274,14 @@ public void setGoogleSearchRetrieval(boolean googleSearchRetrieval) {
 		this.googleSearchRetrieval = googleSearchRetrieval;
 	}
 
+	public List<SafetySetting> getSafetySettings() {
+		return safetySettings;
+	}
+
+	public void setSafetySettings(List<SafetySetting> safetySettings) {
+		this.safetySettings = safetySettings;
+	}
+
 	@Override
 	public Boolean getProxyToolCalls() {
 		return this.proxyToolCalls;
@@ -408,6 +421,11 @@ public Builder withGoogleSearchRetrieval(boolean googleSearch) {
 			return this;
 		}
 
+		public Builder withSafetySettings(List<SafetySetting> safetySettings) {
+			this.options.safetySettings = safetySettings;
+			return this;
+		}
+
 		public Builder withProxyToolCalls(boolean proxyToolCalls) {
 			this.options.proxyToolCalls = proxyToolCalls;
 			return this;

models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModelIT.java

@@ -25,6 +25,8 @@
 
 import com.google.cloud.vertexai.Transport;
 import com.google.cloud.vertexai.VertexAI;
+import com.google.cloud.vertexai.api.HarmCategory;
+import com.google.cloud.vertexai.api.SafetySetting;
 import org.jetbrains.annotations.NotNull;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -91,6 +93,22 @@ void googleSearchTool() {
 		assertThat(response.getResult().getOutput().getContent()).containsAnyOf("Blackbeard", "Bartholomew");
 	}
 
+	@Test
+	void testSafetySettings() {
+		List<SafetySetting> safetySettings = List.of(
+				SafetySetting.newBuilder()
+						.setCategory(HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT)
+						.setThreshold(SafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
+						.build()
+		);
+		Prompt prompt = new Prompt(
+				"What are common digital attack vectors?",
+				VertexAiGeminiChatOptions.builder().withSafetySettings(safetySettings).build()
+		);
+		ChatResponse response = this.chatModel.call(prompt);
+		assertThat(response.getResult().getMetadata().getFinishReason()).isEqualTo("SAFETY");
+	}
+
 	@NotNull
 	private Prompt createPrompt(VertexAiGeminiChatOptions chatOptions) {
 		String request = "Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.";