feat: publish v1.0.0-rc.36-custom.1

relaykit/relayconvert/safety_settings_test.go

@@ -0,0 +1,94 @@
+package relayconvert
+
+import (
+	"context"
+	"testing"
+
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta"
+	kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestOpenAIToGeminiSafetySettings(t *testing.T) {
+	converters := []struct {
+		name    string
+		convert func(t *testing.T, meta convmeta.Meta) *dto.GeminiChatRequest
+	}{
+		{
+			name: "chat completions",
+			convert: func(t *testing.T, meta convmeta.Meta) *dto.GeminiChatRequest {
+				t.Helper()
+				got, err := OpenAIChatRequestToGeminiGenerateContent(context.Background(), dto.GeneralOpenAIRequest{
+					Model: "gemini-test",
+					Messages: []dto.Message{
+						{Role: "user", Content: "hello"},
+					},
+				}, meta)
+				require.NoError(t, err)
+				return got
+			},
+		},
+		{
+			name: "responses",
+			convert: func(t *testing.T, meta convmeta.Meta) *dto.GeminiChatRequest {
+				t.Helper()
+				got, err := OpenAIResponsesRequestToGeminiChat(context.Background(), &dto.OpenAIResponsesRequest{
+					Model: "gemini-test",
+					Input: []byte(`"hello"`),
+				}, meta)
+				require.NoError(t, err)
+				return got
+			},
+		},
+	}
+
+	for _, converter := range converters {
+		t.Run(converter.name, func(t *testing.T) {
+			t.Run("nil meta", func(t *testing.T) {
+				got := converter.convert(t, nil)
+				assert.Empty(t, got.SafetySettings)
+				body, err := kitutil.Marshal(got)
+				require.NoError(t, err)
+				assert.NotContains(t, string(body), `"safetySettings"`)
+			})
+
+			t.Run("zero options", func(t *testing.T) {
+				got := converter.convert(t, &convmeta.Values{})
+				assert.Empty(t, got.SafetySettings)
+				body, err := kitutil.Marshal(got)
+				require.NoError(t, err)
+				assert.NotContains(t, string(body), `"safetySettings"`)
+			})
+
+			t.Run("empty thresholds", func(t *testing.T) {
+				got := converter.convert(t, &convmeta.Values{Options: &convmeta.Options{
+					Gemini: convmeta.GeminiOptions{
+						SafetySetting: func(category string) string {
+							if category == "HARM_CATEGORY_HARASSMENT" {
+								return "BLOCK_NONE"
+							}
+							return ""
+						},
+					},
+				}})
+				assert.Equal(t, []dto.GeminiChatSafetySettings{
+					{Category: "HARM_CATEGORY_HARASSMENT", Threshold: "BLOCK_NONE"},
+				}, got.SafetySettings)
+			})
+
+			t.Run("nonempty thresholds", func(t *testing.T) {
+				got := converter.convert(t, &convmeta.Values{Options: &convmeta.Options{
+					Gemini: convmeta.GeminiOptions{
+						SafetySetting: func(string) string { return "OFF" },
+					},
+				}})
+				require.Len(t, got.SafetySettings, 4)
+				for _, setting := range got.SafetySettings {
+					assert.Equal(t, "OFF", setting.Threshold)
+				}
+			})
+		})
+	}
+}