cneicy
a5f7b8e586
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
302 lines
13 KiB
Go
302 lines
13 KiB
Go
package relay
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
|
"github.com/QuantumNous/new-api/relay/helper"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
kitreasoning "github.com/QuantumNous/new-api/relaykit/relayconvert/reasoning"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/QuantumNous/new-api/setting/model_setting"
|
|
hosttypes "github.com/QuantumNous/new-api/types"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
type imageReservation struct {
|
|
held, limit int
|
|
}
|
|
|
|
func (s *imageReservation) Reserve(quota int) error {
|
|
if quota > s.limit {
|
|
return errors.New("insufficient image quota")
|
|
}
|
|
s.held = max(s.held, quota)
|
|
return nil
|
|
}
|
|
func (s *imageReservation) GetPreConsumedQuota() int { return s.held }
|
|
func (*imageReservation) Settle(int) error { return nil }
|
|
func (*imageReservation) Refund(*gin.Context) {}
|
|
func (*imageReservation) NeedsRefund() bool { return false }
|
|
|
|
func TestImageRequestReservesFinalQuantityBeforeUpstream(t *testing.T) {
|
|
service.InitHttpClient()
|
|
for _, tc := range []struct {
|
|
name, body string
|
|
tiered, passThrough, insufficient, retryToOpenAI bool
|
|
override any
|
|
count, status int
|
|
}{
|
|
{name: "Ali nested count", body: `{"model":"z-image","n":1,"parameters":{"n":4}}`, count: 4},
|
|
{name: "Ali empty parameters", body: `{"model":"z-image","n":2,"parameters":{}}`, count: 2},
|
|
{name: "legacy channel quantity override", body: `{"model":"z-image","n":1}`, override: 4, count: 4},
|
|
{name: "expression channel quantity override", body: `{"model":"z-image","n":1}`, override: 4, count: 4, tiered: true},
|
|
{name: "pass-through quantity", body: `{"model":"z-image","n":2,"parameters":{}}`, count: 2, passThrough: true},
|
|
{name: "zero override rejected", body: `{"model":"z-image","n":1}`, override: 0, status: http.StatusBadRequest},
|
|
{name: "oversized override rejected", body: `{"model":"z-image","n":1}`, override: 129, status: http.StatusBadRequest},
|
|
{name: "insufficient reservation blocks upstream", body: `{"model":"z-image","n":1}`, override: 4, count: 4, insufficient: true, status: http.StatusForbidden},
|
|
{name: "retry drops Ali quantity and surcharge", body: `{"model":"z-image","n":1,"parameters":{"n":4,"prompt_extend":true}}`, count: 1, retryToOpenAI: true},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
received := make(chan []byte, 1)
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
body, err := io.ReadAll(r.Body)
|
|
if err == nil {
|
|
received <- body
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusBadGateway)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"fixture upstream failure","type":"upstream_error"}}`)
|
|
}))
|
|
t.Cleanup(upstream.Close)
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(tc.body))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
channel := constant.ChannelTypeAli
|
|
if tc.retryToOpenAI {
|
|
channel = constant.ChannelTypeOpenAI
|
|
}
|
|
common.SetContextKey(c, constant.ContextKeyChannelType, channel)
|
|
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, upstream.URL)
|
|
common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{PassThroughBodyEnabled: tc.passThrough})
|
|
if tc.override != nil {
|
|
common.SetContextKey(c, constant.ContextKeyChannelParamOverride, map[string]any{"operations": []any{map[string]any{"path": "parameters.n", "mode": "set", "value": tc.override}}})
|
|
}
|
|
request, err := helper.GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
|
|
require.NoError(t, err)
|
|
reservation := &imageReservation{held: 20000, limit: 500000}
|
|
if tc.insufficient {
|
|
reservation.limit = reservation.held
|
|
}
|
|
info := &relaycommon.RelayInfo{Request: request, OriginModelName: "z-image", RelayMode: relayconstant.RelayModeImagesGenerations,
|
|
RequestURLPath: c.Request.URL.Path, Billing: reservation,
|
|
PriceData: hosttypes.PriceData{UsePrice: true, ModelPrice: 0.04, GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1}},
|
|
}
|
|
if tc.tiered {
|
|
expr := `tier("image", fixed(0.04)) * image_count`
|
|
info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expr, ExprHash: billingexpr.ExprHashString(expr), GroupRatio: 1, QuotaPerUnit: common.QuotaPerUnit, EstimatedImageCount: common.GetPointer(1)}
|
|
info.BillingRequestInput = &billingexpr.RequestInput{Body: []byte(tc.body), ImageCount: common.GetPointer(1)}
|
|
info.PriceData.UsePrice = false
|
|
}
|
|
if tc.retryToOpenAI {
|
|
reservation.held = 160000
|
|
info.PriceData.AddOtherRatio("n", 4)
|
|
info.PriceData.AddOtherRatio("prompt_extend", 2)
|
|
info.BillingImageCount = common.GetPointer(3)
|
|
}
|
|
apiErr := ImageHelper(c, info)
|
|
require.NotNil(t, apiErr)
|
|
if tc.status != 0 {
|
|
assert.Equal(t, tc.status, apiErr.StatusCode)
|
|
assert.Empty(t, received, "no upstream submission before quantity validation and reservation")
|
|
return
|
|
}
|
|
assert.Equal(t, http.StatusBadGateway, apiErr.StatusCode)
|
|
require.Len(t, received, 1)
|
|
body := <-received
|
|
path := "parameters.n"
|
|
if tc.retryToOpenAI {
|
|
path = "n"
|
|
}
|
|
assert.Equal(t, int64(tc.count), gjson.GetBytes(body, path).Int())
|
|
assert.Equal(t, tc.count*20000, info.PriceData.QuotaToPreConsume)
|
|
assert.GreaterOrEqual(t, reservation.held, info.PriceData.QuotaToPreConsume)
|
|
assert.Nil(t, info.BillingImageCount)
|
|
if tc.tiered {
|
|
assert.Equal(t, tc.body, string(info.BillingRequestInput.Body))
|
|
assert.Equal(t, 1, *info.BillingRequestInput.ImageCount)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOptInSafeToolLossRejectedAsBadRequestWithAdminDiagnostics(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
|
|
info := &relaycommon.RelayInfo{
|
|
OriginModelName: "gpt-4o",
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
UpstreamModelName: "gpt-4o",
|
|
ChannelOtherSettings: dto.ChannelOtherSettings{
|
|
ToolLossPolicy: string(types.ConversionLossPolicySafe),
|
|
},
|
|
},
|
|
}
|
|
|
|
tools, err := common.Marshal([]map[string]any{{"codeExecution": map[string]any{}}})
|
|
require.NoError(t, err)
|
|
req := &dto.GeminiChatRequest{
|
|
Contents: []dto.GeminiChatContent{
|
|
{Role: "user", Parts: []dto.GeminiPart{{Text: "run this"}}},
|
|
},
|
|
Tools: tools,
|
|
}
|
|
|
|
result, convErr := service.ConvertRequest(c, info, types.RelayFormatOpenAI, req)
|
|
require.Error(t, convErr)
|
|
var loss *types.ConversionLossError
|
|
require.ErrorAs(t, convErr, &loss)
|
|
require.NotEmpty(t, loss.Diagnostics)
|
|
require.NotNil(t, result)
|
|
|
|
apiErr := newConvertRequestFailedError(c, info, convErr)
|
|
require.NotNil(t, apiErr)
|
|
assert.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
|
|
assert.Equal(t, types.ErrorCodeConvertRequestFailed, apiErr.GetErrorCode())
|
|
assert.True(t, types.IsSkipRetryError(apiErr))
|
|
|
|
diagnostics := info.ConversionDiagnostics()
|
|
require.NotEmpty(t, diagnostics)
|
|
assert.True(t, hasHostDiagnosticCode(diagnostics, "unsupported_hosted_tool"))
|
|
|
|
other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1)
|
|
adminInfo, ok := other.Snapshot()["admin_info"].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Contains(t, adminInfo, "conversion_diagnostics")
|
|
}
|
|
|
|
func TestUnknownModelModifierIsBadRequestWithoutRetry(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
|
|
info := &relaycommon.RelayInfo{
|
|
OriginModelName: "m@thinkin:on",
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
UpstreamModelName: "m@thinkin:on",
|
|
},
|
|
}
|
|
err := helper.ApplyReasoningModelSuffix(c, info)
|
|
require.Error(t, err)
|
|
require.True(t, kitreasoning.IsClientError(err))
|
|
assert.Contains(t, err.Error(), `unsupported model modifier "thinkin"`)
|
|
assert.Contains(t, err.Error(), "re:")
|
|
|
|
apiErr := newConvertRequestFailedError(c, info, err)
|
|
require.NotNil(t, apiErr)
|
|
assert.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
|
|
assert.Equal(t, types.ErrorCodeConvertRequestFailed, apiErr.GetErrorCode())
|
|
assert.True(t, types.IsSkipRetryError(apiErr))
|
|
}
|
|
|
|
func hasHostDiagnosticCode(diagnostics []types.ConversionDiagnostic, code string) bool {
|
|
for _, diagnostic := range diagnostics {
|
|
if diagnostic.Code == code {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func TestGeminiThinkingControlsConvertBestEffort(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
settings := model_setting.GetGeminiSettings()
|
|
originalAdapter := settings.ThinkingAdapterEnabled
|
|
t.Cleanup(func() { settings.ThinkingAdapterEnabled = originalAdapter })
|
|
|
|
t.Run("openai thinking_budget on gemini 3 becomes thinkingLevel", func(t *testing.T) {
|
|
settings.ThinkingAdapterEnabled = false
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
req := &dto.GeneralOpenAIRequest{
|
|
Model: "gemini-3.1-pro-preview-thinking",
|
|
Messages: []dto.Message{{Role: "user", Content: "hello"}},
|
|
ExtraBody: []byte(`{"google":{"thinking_config":{"thinking_budget":8192}}}`),
|
|
}
|
|
info := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "gemini-3.1-pro-preview-thinking",
|
|
Request: req,
|
|
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-3.1-pro-preview"},
|
|
}
|
|
require.NoError(t, helper.ApplyReasoningModelSuffix(c, info, req))
|
|
|
|
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, req)
|
|
require.NoError(t, err)
|
|
converted, ok := result.Value.(*dto.GeminiChatRequest)
|
|
require.True(t, ok)
|
|
require.NotNil(t, converted.GenerationConfig.ThinkingConfig)
|
|
assert.Equal(t, "medium", converted.GenerationConfig.ThinkingConfig.ThinkingLevel)
|
|
assert.Nil(t, converted.GenerationConfig.ThinkingConfig.ThinkingBudget)
|
|
assert.True(t, hasHostDiagnosticCode(info.ConversionDiagnostics(), "gemini_budget_to_level"))
|
|
})
|
|
|
|
t.Run("thinking alias with native thinking_level keeps the level", func(t *testing.T) {
|
|
settings.ThinkingAdapterEnabled = true
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
req := &dto.GeneralOpenAIRequest{
|
|
Model: "gemini-3.1-pro-preview-thinking",
|
|
Messages: []dto.Message{{Role: "user", Content: "hello"}},
|
|
ExtraBody: []byte(`{"google":{"thinking_config":{"thinking_level":"low"}}}`),
|
|
}
|
|
info := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "gemini-3.1-pro-preview-thinking",
|
|
Request: req,
|
|
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-3.1-pro-preview-thinking"},
|
|
}
|
|
require.NoError(t, helper.ApplyReasoningModelSuffix(c, info, req))
|
|
assert.Equal(t, "gemini-3.1-pro-preview", info.UpstreamModelName)
|
|
|
|
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, req)
|
|
require.NoError(t, err)
|
|
converted, ok := result.Value.(*dto.GeminiChatRequest)
|
|
require.True(t, ok)
|
|
require.NotNil(t, converted.GenerationConfig.ThinkingConfig)
|
|
assert.Equal(t, "low", converted.GenerationConfig.ThinkingConfig.ThinkingLevel)
|
|
assert.Nil(t, converted.GenerationConfig.ThinkingConfig.ThinkingBudget)
|
|
})
|
|
|
|
t.Run("gemini thinkingBudget converts to openai reasoning_effort", func(t *testing.T) {
|
|
settings.ThinkingAdapterEnabled = false
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3.1-pro-preview:generateContent", nil)
|
|
budget := 8192
|
|
req := &dto.GeminiChatRequest{
|
|
Contents: []dto.GeminiChatContent{{Role: "user", Parts: []dto.GeminiPart{{Text: "hello"}}}},
|
|
GenerationConfig: dto.GeminiChatGenerationConfig{
|
|
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingBudget: &budget},
|
|
},
|
|
}
|
|
info := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatGemini,
|
|
OriginModelName: "gemini-3.1-pro-preview",
|
|
Request: req,
|
|
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-5"},
|
|
}
|
|
|
|
result, err := service.ConvertRequest(c, info, types.RelayFormatOpenAI, req)
|
|
require.NoError(t, err)
|
|
converted, ok := result.Value.(*dto.GeneralOpenAIRequest)
|
|
require.True(t, ok)
|
|
assert.Equal(t, "medium", converted.ReasoningEffort)
|
|
assert.True(t, hasHostDiagnosticCode(info.ConversionDiagnostics(), "gemini_budget_to_level"))
|
|
})
|
|
}
|