Skip to content

Commit f4eb375

Browse files
markpollackspring-builds
markpollack
authored and
spring-builds
committed
feat(vertex-gemini): Add safety ratings to response metadata
Extract safety ratings from Gemini response candidates and include them in AssistantMessage metadata. Update Google Cloud BOM to 26.72.0. Fixes #687 (cherry picked from commit 2710cab)
1 parent a7e3752 commit f4eb375

File tree

5 files changed

+230
-2
lines changed

5 files changed

+230
-2
lines changed

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

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@
8787
import org.springframework.ai.tool.definition.ToolDefinition;
8888
import org.springframework.ai.vertexai.gemini.api.VertexAiGeminiApi;
8989
import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiConstants;
90+
import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiSafetyRating;
9091
import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiSafetySetting;
9192
import org.springframework.ai.vertexai.gemini.schema.VertexAiSchemaConverter;
9293
import org.springframework.ai.vertexai.gemini.schema.VertexToolCallingManager;
@@ -594,8 +595,16 @@ protected List<Generation> responseCandidateToGeneration(Candidate candidate) {
594595
VertexAiGeminiApi.LogProbs logprobs = new VertexAiGeminiApi.LogProbs(candidate.getAvgLogprobs(), topCandidates,
595596
chosenCandidates);
596597

598+
// Extract safety ratings from the candidate
599+
List<VertexAiGeminiSafetyRating> safetyRatings = candidate.getSafetyRatingsList()
600+
.stream()
601+
.map(sr -> new VertexAiGeminiSafetyRating(toSafetyRatingHarmCategory(sr.getCategory()),
602+
toSafetyRatingHarmProbability(sr.getProbability()), sr.getBlocked(), sr.getProbabilityScore(),
603+
toSafetyRatingHarmSeverity(sr.getSeverity()), sr.getSeverityScore()))
604+
.toList();
605+
597606
Map<String, Object> messageMetadata = Map.of("candidateIndex", candidateIndex, "finishReason",
598-
candidateFinishReason, "logprobs", logprobs);
607+
candidateFinishReason, "logprobs", logprobs, "safetyRatings", safetyRatings);
599608

600609
ChatGenerationMetadata chatGenerationMetadata = ChatGenerationMetadata.builder()
601610
.finishReason(candidateFinishReason.name())
@@ -633,6 +642,42 @@ private DefaultUsage getDefaultUsage(GenerateContentResponse.UsageMetadata usage
633642
usageMetadata.getTotalTokenCount(), usageMetadata);
634643
}
635644

645+
private VertexAiGeminiSafetyRating.HarmCategory toSafetyRatingHarmCategory(
646+
com.google.cloud.vertexai.api.HarmCategory category) {
647+
return switch (category) {
648+
case HARM_CATEGORY_HATE_SPEECH -> VertexAiGeminiSafetyRating.HarmCategory.HARM_CATEGORY_HATE_SPEECH;
649+
case HARM_CATEGORY_DANGEROUS_CONTENT ->
650+
VertexAiGeminiSafetyRating.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT;
651+
case HARM_CATEGORY_HARASSMENT -> VertexAiGeminiSafetyRating.HarmCategory.HARM_CATEGORY_HARASSMENT;
652+
case HARM_CATEGORY_SEXUALLY_EXPLICIT ->
653+
VertexAiGeminiSafetyRating.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT;
654+
case HARM_CATEGORY_CIVIC_INTEGRITY -> VertexAiGeminiSafetyRating.HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY;
655+
default -> VertexAiGeminiSafetyRating.HarmCategory.HARM_CATEGORY_UNSPECIFIED;
656+
};
657+
}
658+
659+
private VertexAiGeminiSafetyRating.HarmProbability toSafetyRatingHarmProbability(
660+
com.google.cloud.vertexai.api.SafetyRating.HarmProbability probability) {
661+
return switch (probability) {
662+
case NEGLIGIBLE -> VertexAiGeminiSafetyRating.HarmProbability.NEGLIGIBLE;
663+
case LOW -> VertexAiGeminiSafetyRating.HarmProbability.LOW;
664+
case MEDIUM -> VertexAiGeminiSafetyRating.HarmProbability.MEDIUM;
665+
case HIGH -> VertexAiGeminiSafetyRating.HarmProbability.HIGH;
666+
default -> VertexAiGeminiSafetyRating.HarmProbability.HARM_PROBABILITY_UNSPECIFIED;
667+
};
668+
}
669+
670+
private VertexAiGeminiSafetyRating.HarmSeverity toSafetyRatingHarmSeverity(
671+
com.google.cloud.vertexai.api.SafetyRating.HarmSeverity severity) {
672+
return switch (severity) {
673+
case HARM_SEVERITY_NEGLIGIBLE -> VertexAiGeminiSafetyRating.HarmSeverity.HARM_SEVERITY_NEGLIGIBLE;
674+
case HARM_SEVERITY_LOW -> VertexAiGeminiSafetyRating.HarmSeverity.HARM_SEVERITY_LOW;
675+
case HARM_SEVERITY_MEDIUM -> VertexAiGeminiSafetyRating.HarmSeverity.HARM_SEVERITY_MEDIUM;
676+
case HARM_SEVERITY_HIGH -> VertexAiGeminiSafetyRating.HarmSeverity.HARM_SEVERITY_HIGH;
677+
default -> VertexAiGeminiSafetyRating.HarmSeverity.HARM_SEVERITY_UNSPECIFIED;
678+
};
679+
}
680+
636681
private VertexAiGeminiChatOptions vertexAiGeminiChatOptions(Prompt prompt) {
637682
VertexAiGeminiChatOptions updatedRuntimeOptions = VertexAiGeminiChatOptions.builder().build();
638683
if (prompt.getOptions() != null) {

‎models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/common/VertexAiGeminiSafetyRating.java‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Copyright 2024-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.vertexai.gemini.common;
18+
19+
/**
20+
* Represents a safety rating returned by the Vertex AI Gemini API for generated content.
21+
* Safety ratings indicate the probability and severity of harmful content in a specific
22+
* category.
23+
*
24+
* @author Mark Pollack
25+
* @since 1.1.1
26+
* @see VertexAiGeminiSafetySetting
27+
*/
28+
public record VertexAiGeminiSafetyRating(HarmCategory category, HarmProbability probability, boolean blocked,
29+
float probabilityScore, HarmSeverity severity, float severityScore) {
30+
31+
/**
32+
* Enum representing different categories of harmful content.
33+
*/
34+
public enum HarmCategory {
35+
36+
HARM_CATEGORY_UNSPECIFIED, HARM_CATEGORY_HATE_SPEECH, HARM_CATEGORY_DANGEROUS_CONTENT, HARM_CATEGORY_HARASSMENT,
37+
HARM_CATEGORY_SEXUALLY_EXPLICIT, HARM_CATEGORY_CIVIC_INTEGRITY
38+
39+
}
40+
41+
/**
42+
* Enum representing the probability levels of harmful content.
43+
*/
44+
public enum HarmProbability {
45+
46+
HARM_PROBABILITY_UNSPECIFIED, NEGLIGIBLE, LOW, MEDIUM, HIGH
47+
48+
}
49+
50+
/**
51+
* Enum representing the severity levels of harmful content.
52+
*/
53+
public enum HarmSeverity {
54+
55+
HARM_SEVERITY_UNSPECIFIED, HARM_SEVERITY_NEGLIGIBLE, HARM_SEVERITY_LOW, HARM_SEVERITY_MEDIUM, HARM_SEVERITY_HIGH
56+
57+
}
58+
59+
}

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
import org.springframework.ai.tool.annotation.Tool;
4848
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatModel.ChatModel;
4949
import org.springframework.ai.vertexai.gemini.api.VertexAiGeminiApi;
50+
import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiSafetyRating;
5051
import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiSafetySetting;
5152
import org.springframework.beans.factory.annotation.Autowired;
5253
import org.springframework.beans.factory.annotation.Value;
@@ -247,6 +248,50 @@ void logprobs() {
247248
assertThat(logprobs.chosenCandidates()).isNotEmpty();
248249
}
249250

251+
@Test
252+
@SuppressWarnings("unchecked")
253+
void safetyRatingsMetadataIsPresent() {
254+
// Use safety settings with BLOCK_LOW_AND_ABOVE to ensure safety evaluation occurs
255+
// and ratings are returned (similar to Python SDK example)
256+
List<VertexAiGeminiSafetySetting> safetySettings = List.of(
257+
VertexAiGeminiSafetySetting.builder()
258+
.withCategory(VertexAiGeminiSafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT)
259+
.withThreshold(VertexAiGeminiSafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
260+
.build(),
261+
VertexAiGeminiSafetySetting.builder()
262+
.withCategory(VertexAiGeminiSafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH)
263+
.withThreshold(VertexAiGeminiSafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
264+
.build(),
265+
VertexAiGeminiSafetySetting.builder()
266+
.withCategory(VertexAiGeminiSafetySetting.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT)
267+
.withThreshold(VertexAiGeminiSafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
268+
.build(),
269+
VertexAiGeminiSafetySetting.builder()
270+
.withCategory(VertexAiGeminiSafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT)
271+
.withThreshold(VertexAiGeminiSafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
272+
.build());
273+
274+
// Use a prompt that should trigger safety evaluation
275+
String prompt = "Write a list of 5 disrespectful things that I might say to the universe after stubbing my toe in the dark:";
276+
277+
ChatResponse response = this.chatModel
278+
.call(new Prompt(prompt, VertexAiGeminiChatOptions.builder().safetySettings(safetySettings).build()));
279+
280+
// Safety ratings should be present in the AssistantMessage metadata
281+
var safetyRatings = (List<VertexAiGeminiSafetyRating>) response.getResult()
282+
.getOutput()
283+
.getMetadata()
284+
.get("safetyRatings");
285+
286+
assertThat(safetyRatings).isNotNull();
287+
assertThat(safetyRatings).isNotEmpty();
288+
289+
// Verify safety rating structure
290+
VertexAiGeminiSafetyRating firstRating = safetyRatings.get(0);
291+
assertThat(firstRating.category()).isNotNull();
292+
assertThat(firstRating.probability()).isNotNull();
293+
}
294+
250295
@Test
251296
void beanStreamOutputConverterRecords() {
252297

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@
286286
<djl.version>0.32.0</djl.version>
287287
<onnxruntime.version>1.19.2</onnxruntime.version>
288288
<oci-sdk-version>3.63.1</oci-sdk-version>
289-
<com.google.cloud.version>26.60.0</com.google.cloud.version>
289+
<com.google.cloud.version>26.72.0</com.google.cloud.version>
290290
<com.google.genai.version>1.28.0</com.google.genai.version>
291291
<ibm.sdk.version>9.20.0</ibm.sdk.version>
292292
<jsonschema.version>4.38.0</jsonschema.version>

‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/vertexai-gemini-chat.adoc‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,85 @@ var response = this.chatModel.call(new Prompt(List.of(userMessage)));
215215
----
216216

217217

218+
== Safety Settings and Safety Ratings
219+
220+
The Vertex AI Gemini API provides safety filtering capabilities to help you control harmful content in both prompts and responses.
221+
For more details, see the https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters[Vertex AI Safety Filters documentation].
222+
223+
=== Configuring Safety Settings
224+
225+
You can configure safety settings to control the threshold at which content is blocked for different harm categories.
226+
The available harm categories are:
227+
228+
* `HARM_CATEGORY_HATE_SPEECH` - Hate speech content
229+
* `HARM_CATEGORY_DANGEROUS_CONTENT` - Dangerous content
230+
* `HARM_CATEGORY_HARASSMENT` - Harassment content
231+
* `HARM_CATEGORY_SEXUALLY_EXPLICIT` - Sexually explicit content
232+
* `HARM_CATEGORY_CIVIC_INTEGRITY` - Civic integrity content
233+
234+
The available threshold levels are:
235+
236+
* `BLOCK_LOW_AND_ABOVE` - Block when low, medium, or high probability of unsafe content
237+
* `BLOCK_MEDIUM_AND_ABOVE` - Block when medium or high probability of unsafe content
238+
* `BLOCK_ONLY_HIGH` - Block only when high probability of unsafe content
239+
* `BLOCK_NONE` - Never block (use with caution)
240+
241+
[source,java]
242+
----
243+
List<VertexAiGeminiSafetySetting> safetySettings = List.of(
244+
VertexAiGeminiSafetySetting.builder()
245+
.withCategory(VertexAiGeminiSafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT)
246+
.withThreshold(VertexAiGeminiSafetySetting.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
247+
.build(),
248+
VertexAiGeminiSafetySetting.builder()
249+
.withCategory(VertexAiGeminiSafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH)
250+
.withThreshold(VertexAiGeminiSafetySetting.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE)
251+
.build());
252+
253+
ChatResponse response = chatModel.call(new Prompt("Your prompt here",
254+
VertexAiGeminiChatOptions.builder()
255+
.safetySettings(safetySettings)
256+
.build()));
257+
----
258+
259+
=== Accessing Safety Ratings in Responses
260+
261+
When safety settings are configured, the Gemini API returns safety ratings for each response candidate.
262+
These ratings indicate the probability and severity of harmful content in each category.
263+
264+
Safety ratings are available in the `AssistantMessage` metadata under the key `"safetyRatings"`:
265+
266+
[source,java]
267+
----
268+
ChatResponse response = chatModel.call(new Prompt(prompt,
269+
VertexAiGeminiChatOptions.builder()
270+
.safetySettings(safetySettings)
271+
.build()));
272+
273+
// Access safety ratings from the response
274+
List<VertexAiGeminiSafetyRating> safetyRatings =
275+
(List<VertexAiGeminiSafetyRating>) response.getResult()
276+
.getOutput()
277+
.getMetadata()
278+
.get("safetyRatings");
279+
280+
for (VertexAiGeminiSafetyRating rating : safetyRatings) {
281+
System.out.println("Category: " + rating.category());
282+
System.out.println("Probability: " + rating.probability());
283+
System.out.println("Severity: " + rating.severity());
284+
System.out.println("Blocked: " + rating.blocked());
285+
}
286+
----
287+
288+
The `VertexAiGeminiSafetyRating` record contains:
289+
290+
* `category` - The harm category (e.g., `HARM_CATEGORY_HARASSMENT`)
291+
* `probability` - The probability level (`NEGLIGIBLE`, `LOW`, `MEDIUM`, `HIGH`)
292+
* `blocked` - Whether the content was blocked due to this rating
293+
* `probabilityScore` - The raw probability score (0.0 to 1.0)
294+
* `severity` - The severity level (`HARM_SEVERITY_NEGLIGIBLE`, `HARM_SEVERITY_LOW`, `HARM_SEVERITY_MEDIUM`, `HARM_SEVERITY_HIGH`)
295+
* `severityScore` - The raw severity score (0.0 to 1.0)
296+
218297
== Sample Controller
219298

220299
https://start.spring.io/[Create] a new Spring Boot project and add the `spring-ai-starter-model-vertex-ai-gemini` to your pom (or gradle) dependencies.

0 commit comments

Comments
 (0)