Skip to content

Commit 32ce458

Browse files
sukumargaonkaraabchooyuzisunmathetake
authored and
nutanix-Hrushikesh
committed
feat: add gemini safety ratings to ChatCompletion responses (envoyproxy#1287)
**Description** This PR adds support for safety ratings in ChatCompletion responses when using GCP Vertex AI. The safety ratings are copied from the Vertex AI response as-is and included in the OpenAI-compatible response format. The implementation adds a new `SafetyRatings` field to the `ChatCompletionResponseChoiceMessage` struct. GCP Safety Ratings doc: [1] Key changes: - Added `SafetyRatings` field to OpenAI API schema for chat completion responses - Updated Gemini translator to map safety ratings from Vertex AI responses - Safety ratings are only included when present in the backend response **Special notes for reviewers (if applicable)** The safety ratings follow the GCP Vertex AI format and are passed through unchanged to maintain compatibility with Google's safety rating system. The field is optional and only populated when safety ratings are present in the upstream response. [1]: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/GenerateContentResponse#SafetyRating --------- Signed-off-by: Sukumar Gaonkar <[email protected]> Signed-off-by: Dan Sun <[email protected]> Co-authored-by: Aaron Choo <[email protected]> Co-authored-by: Dan Sun <[email protected]> Co-authored-by: Takeshi Yoneda <[email protected]> Signed-off-by: Hrushikesh Patil <[email protected]>
1 parent 0344926 commit 32ce458

File tree

4 files changed

+169
-0
lines changed

4 files changed

+169
-0
lines changed

‎internal/apischema/openai/openai.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1220,6 +1220,13 @@ type ChatCompletionResponseChoiceMessage struct {
12201220
// ReasoningContent is used to hold any non-standard fields from the backend which supports reasoning,
12211221
// like "reasoningContent" from AWS Bedrock.
12221222
ReasoningContent *ReasoningContentUnion `json:"reasoning_content,omitempty"`
1223+
1224+
// GCPVertexAI specific fields.
1225+
1226+
// SafetyRatings contains safety ratings copied from the GCP Vertex AI response as-is.
1227+
// List of ratings for the safety of a response candidate. There is at most one rating per category.
1228+
// https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/GenerateContentResponse#SafetyRating
1229+
SafetyRatings []*genai.SafetyRating `json:"safety_ratings,omitempty"`
12231230
}
12241231

12251232
// URLCitation contains citation information for web search results.

‎internal/apischema/openai/openai_test.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/openai/openai-go/v2"
1616
"github.com/openai/openai-go/v2/packages/param"
1717
"github.com/stretchr/testify/require"
18+
"google.golang.org/genai"
1819
"k8s.io/utils/ptr"
1920
)
2021

@@ -992,6 +993,69 @@ func TestChatCompletionResponse(t *testing.T) {
992993
}
993994
}`,
994995
},
996+
{
997+
name: "response with safety settings",
998+
response: ChatCompletionResponse{
999+
ID: "chatcmpl-safety-test",
1000+
Created: JSONUNIXTime(time.Unix(1755135425, 0)),
1001+
Model: "gpt-4.1-nano",
1002+
Object: "chat.completion",
1003+
Choices: []ChatCompletionResponseChoice{
1004+
{
1005+
Index: 0,
1006+
FinishReason: ChatCompletionChoicesFinishReasonStop,
1007+
Message: ChatCompletionResponseChoiceMessage{
1008+
Role: "assistant",
1009+
Content: ptr.To("This is a safe response"),
1010+
SafetyRatings: []*genai.SafetyRating{
1011+
{
1012+
Category: genai.HarmCategoryHarassment,
1013+
Probability: genai.HarmProbabilityLow,
1014+
},
1015+
{
1016+
Category: genai.HarmCategorySexuallyExplicit,
1017+
Probability: genai.HarmProbabilityNegligible,
1018+
},
1019+
},
1020+
},
1021+
},
1022+
},
1023+
Usage: ChatCompletionResponseUsage{
1024+
CompletionTokens: 5,
1025+
PromptTokens: 3,
1026+
TotalTokens: 8,
1027+
},
1028+
},
1029+
expected: `{
1030+
"id": "chatcmpl-safety-test",
1031+
"object": "chat.completion",
1032+
"created": 1755135425,
1033+
"model": "gpt-4.1-nano",
1034+
"choices": [{
1035+
"index": 0,
1036+
"message": {
1037+
"role": "assistant",
1038+
"content": "This is a safe response",
1039+
"safety_ratings": [
1040+
{
1041+
"category": "HARM_CATEGORY_HARASSMENT",
1042+
"probability": "LOW"
1043+
},
1044+
{
1045+
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
1046+
"probability": "NEGLIGIBLE"
1047+
}
1048+
]
1049+
},
1050+
"finish_reason": "stop"
1051+
}],
1052+
"usage": {
1053+
"prompt_tokens": 3,
1054+
"completion_tokens": 5,
1055+
"total_tokens": 8
1056+
}
1057+
}`,
1058+
},
9951059
}
9961060

9971061
for _, tc := range testCases {

‎internal/extproc/translator/gemini_helper.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -486,6 +486,14 @@ func geminiCandidatesToOpenAIChoices(candidates []*genai.Candidate) ([]openai.Ch
486486
choice.Message = message
487487
}
488488

489+
if candidate.SafetyRatings != nil {
490+
if choice.Message.Role == "" {
491+
choice.Message.Role = openai.ChatMessageRoleAssistant
492+
}
493+
494+
choice.Message.SafetyRatings = candidate.SafetyRatings
495+
}
496+
489497
// Handle logprobs if available.
490498
if candidate.LogprobsResult != nil {
491499
choice.Logprobs = geminiLogprobsToOpenAILogprobs(*candidate.LogprobsResult)

‎internal/extproc/translator/openai_gcpvertexai_test.go‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -721,6 +721,96 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_ResponseBody(t *testing.T
721721
TotalTokens: 25,
722722
},
723723
},
724+
{
725+
name: "response with safety ratings",
726+
respHeaders: map[string]string{
727+
"content-type": "application/json",
728+
},
729+
body: `{
730+
"candidates": [
731+
{
732+
"content": {
733+
"parts": [
734+
{
735+
"text": "This is a safe response from the AI assistant."
736+
}
737+
]
738+
},
739+
"finishReason": "STOP",
740+
"safetyRatings": [
741+
{
742+
"category": "HARM_CATEGORY_HARASSMENT",
743+
"probability": "LOW"
744+
},
745+
{
746+
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
747+
"probability": "NEGLIGIBLE"
748+
},
749+
{
750+
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
751+
"probability": "MEDIUM"
752+
}
753+
]
754+
}
755+
],
756+
"promptFeedback": {
757+
"safetyRatings": []
758+
},
759+
"usageMetadata": {
760+
"promptTokenCount": 8,
761+
"candidatesTokenCount": 12,
762+
"totalTokenCount": 20
763+
}
764+
}`,
765+
endOfStream: true,
766+
wantError: false,
767+
wantHeaderMut: &extprocv3.HeaderMutation{
768+
SetHeaders: []*corev3.HeaderValueOption{{
769+
Header: &corev3.HeaderValue{Key: "Content-Length", RawValue: []byte("457")},
770+
}},
771+
},
772+
wantBodyMut: &extprocv3.BodyMutation{
773+
Mutation: &extprocv3.BodyMutation_Body{
774+
Body: []byte(`{
775+
"choices": [
776+
{
777+
"finish_reason": "stop",
778+
"index": 0,
779+
"message": {
780+
"content": "This is a safe response from the AI assistant.",
781+
"role": "assistant",
782+
"safety_ratings": [
783+
{
784+
"category": "HARM_CATEGORY_HARASSMENT",
785+
"probability": "LOW"
786+
},
787+
{
788+
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
789+
"probability": "NEGLIGIBLE"
790+
},
791+
{
792+
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
793+
"probability": "MEDIUM"
794+
}
795+
]
796+
}
797+
}
798+
],
799+
"object": "chat.completion",
800+
"usage": {
801+
"completion_tokens": 12,
802+
"prompt_tokens": 8,
803+
"total_tokens": 20
804+
}
805+
}`),
806+
},
807+
},
808+
wantTokenUsage: LLMTokenUsage{
809+
InputTokens: 8,
810+
OutputTokens: 12,
811+
TotalTokens: 20,
812+
},
813+
},
724814
{
725815
name: "empty response",
726816
respHeaders: map[string]string{

0 commit comments

Comments
 (0)