feat: publish sanitized custom.97 source

This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207

No files matched your search

+2 -11
View File
@@ -5,7 +5,6 @@ import (
"strings"
"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"
"github.com/QuantumNous/new-api/relaykit/dto"
@@ -42,19 +41,11 @@ func ResolveImageBillingRequestInput(c *gin.Context, info *relaycommon.RelayInfo
if !ok {
return input, nil
}
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
if info.ChannelMeta != nil {
channelType = info.ChannelType
}
count, err := request.ImageCount(channelType == constant.ChannelTypeAli)
count, err := request.ImageCount(false)
if err != nil {
return input, err
}
topLevelCount, err := request.ImageCount(false)
if err != nil {
return input, err
}
body := map[string]any{"model": request.Model, "n": topLevelCount, "size": request.Size, "quality": request.Quality}
body := map[string]any{"model": request.Model, "n": count, "size": request.Size, "quality": request.Quality}
if request.BillingParameters != nil {
body["parameters"] = request.BillingParameters
}
+13 -4
View File
@@ -7,6 +7,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
@@ -65,7 +66,7 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error {
jsonData, err := common.Marshal(resp)
if err != nil {
common.SysError("error marshalling stream response: " + err.Error())
common.SysError(common.LogText("error marshalling stream response: %s", err.Error()))
} else {
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)})
@@ -91,7 +92,15 @@ func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data st
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)})
return FlushWriter(c)
err := FlushWriter(c)
if err == nil {
if value, exists := c.Get(relaycommon.UpstreamTimingContextKey); exists {
if timing, ok := value.(*relaycommon.UpstreamTiming); ok {
timing.Mark("first_downstream_flush")
}
}
}
return err
}
func StringData(c *gin.Context, str string) error {
@@ -139,7 +148,7 @@ func Done(c *gin.Context) {
func WssString(c *gin.Context, ws *websocket.Conn, str string) error {
if ws == nil {
logger.LogError(c, "websocket connection is nil")
logger.LogError(c, common.LogText("websocket connection is nil"))
return errors.New("websocket connection is nil")
}
//common.LogInfo(c, fmt.Sprintf("sending message: %s", str))
@@ -152,7 +161,7 @@ func WssObject(c *gin.Context, ws *websocket.Conn, object any) error {
return fmt.Errorf("error marshalling object: %w", err)
}
if ws == nil {
logger.LogError(c, "websocket connection is nil")
logger.LogError(c, common.LogText("websocket connection is nil"))
return errors.New("websocket connection is nil")
}
//common.LogInfo(c, fmt.Sprintf("sending message: %s", jsonData))
+4 -4
View File
@@ -19,16 +19,16 @@ func TestGroupModelPricingPreservesOtherGroupsAndSnapshot(t *testing.T) {
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{"gpt-5.6-sol":"tiered_expr"}`,
"billing_setting.billing_expr": `{"gpt-5.6-sol":"tier(\"standard\", p * 4 + c * 20 + cr * 0.4)"}`,
"billing_setting.group_billing_expr": `{"example-premium":{"gpt-5.6-sol":"tier(\"standard\", p * 5 + c * 30 + cr * 0.5)"}}`,
"group_ratio_setting.group_ratio": `{"example-premium":0.25,"example-standard":0.5}`,
"billing_setting.group_billing_expr": `{"gpt-v-pro":{"gpt-5.6-sol":"tier(\"standard\", p * 5 + c * 30 + cr * 0.5)"}}`,
"group_ratio_setting.group_ratio": `{"gpt-v-pro":0.2052,"gpt-v-mix":0.18}`,
}))
for _, tc := range []struct {
group string
expected int
cost float64
}{
{"example-premium", 1000, 8500},
{"example-standard", 1500, 6400},
{"gpt-v-pro", 513, 8500},
{"gpt-v-mix", 360, 6400},
} {
t.Run(tc.group, func(t *testing.T) {
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
@@ -12,9 +12,9 @@ func TestImageGenerationResolutionMapping(t *testing.T) {
cfg := model_setting.GetGlobalSettings()
oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "example-images", []string{"gpt-image-2", "gpt-image-2-4k"}
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "gpt-images", []string{"gpt-image-2", "gpt-image-2-4k"}
for _, tc := range []struct{ size, tier string }{{"", "1k"}, {"auto", "1k"}, {"1024x1536", "1k"}, {"1536x1024", "1k"}, {"2048x1152", "2k"}, {"2048x2048", "2k"}, {"3840x2160", "4k"}, {"2160x3840", "4k"}, {"4k", "4k"}} {
info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "example-images", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "gpt-images", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
image := &dto.ImageRequest{Model: "gpt-image-2", Size: tc.size}
require.NoError(t, mapImageGenerationResolution(info, image))
require.Equal(t, "gpt-image-2-"+tc.tier, info.UpstreamModelName)
+24 -20
View File
@@ -115,38 +115,42 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
})
}
func TestImageBillingRequestUsesValidatedProviderCount(t *testing.T) {
// Provider `parameters.n` is still bounded at ingress so a malformed multiplier
// is a client error, but the billing quantity always follows the top-level n:
// Ali image requests are served by the alibaba task plugin, which owns the
// DashScope parameters.n semantics.
func TestImageBillingRequestValidatesProviderCountWithoutOverriding(t *testing.T) {
for _, tc := range []struct {
body string
channel, count int
invalid bool
body string
count int
invalid bool
}{
{`{"model":"z-image","n":2,"parameters":{"n":3,"prompt_extend":true}}`, constant.ChannelTypeAli, 3, false},
{`{"model":"gpt-image-2","n":2,"parameters":{"n":3}}`, constant.ChannelTypeOpenAI, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":129}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":-1}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","n":2,"parameters":{}}`, constant.ChannelTypeAli, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":null}}`, constant.ChannelTypeAli, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":0}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":1.5}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":18446744073686646784}}`, constant.ChannelTypeAli, 0, true},
{`{"model":"z-image","parameters":{"n":128}}`, constant.ChannelTypeAli, 128, false},
{`{"model":"z-image","n":0}`, constant.ChannelTypeAli, 1, false},
{`{"model":"gpt-image-2","n":2,"parameters":{"n":0}}`, constant.ChannelTypeOpenAI, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":3,"prompt_extend":true}}`, 2, false},
{`{"model":"gpt-image-2","n":2,"parameters":{"n":3}}`, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":129}}`, 0, true},
{`{"model":"z-image","parameters":{"n":-1}}`, 0, true},
{`{"model":"z-image","n":2,"parameters":{}}`, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":null}}`, 2, false},
{`{"model":"z-image","n":2,"parameters":{"n":0}}`, 2, false},
{`{"model":"z-image","parameters":{"n":1.5}}`, 0, true},
{`{"model":"z-image","parameters":{"n":18446744073686646784}}`, 0, true},
{`{"model":"z-image","parameters":{"n":128}}`, 1, false},
{`{"model":"z-image","n":0}`, 1, false},
{`{"model":"gpt-image-2","n":2,"parameters":{"n":0}}`, 2, false},
} {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewBufferString(tc.body))
c.Request.Header.Set("Content-Type", "application/json")
common.SetContextKey(c, constant.ContextKeyChannelType, tc.channel)
common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeOpenAI)
request, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
if tc.invalid {
require.Error(t, err)
require.Error(t, err, tc.body)
continue
}
require.NoError(t, err)
require.NoError(t, err, tc.body)
input, err := ResolveImageBillingRequestInput(c, &relaycommon.RelayInfo{Request: request}, billingexpr.RequestInput{})
require.NoError(t, err)
require.Equal(t, tc.count, *input.ImageCount)
require.Equal(t, tc.count, *input.ImageCount, tc.body)
cost, _, err := billingexpr.RunExprWithRequest(`tier("image", fixed(0.04)) * image_count`, billingexpr.TokenParams{}, input)
require.NoError(t, err)
require.Equal(t, float64(tc.count)*40000, cost)
+55 -45
View File
@@ -1,11 +1,12 @@
package helper
import (
"errors"
"fmt"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
@@ -14,6 +15,7 @@ import (
"github.com/QuantumNous/new-api/relaykit/relayconvert/reasoning"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
hostreasoning "github.com/QuantumNous/new-api/setting/reasoning"
@@ -22,29 +24,17 @@ import (
"github.com/gin-gonic/gin"
)
func modelPriceNotConfiguredError(modelName string, userId int) error {
func modelPriceNotConfiguredError(c *gin.Context, modelName string, userId int) error {
key := i18n.MsgRelayModelPriceNotConfigured
if model.IsAdmin(userId) {
return fmt.Errorf(
"模型 %s 的价格未配置。请前往「系统设置 → 运营设置」开启自用模式,或在「系统设置 → 分组与模型定价设置」中为该模型配置价格;"+
"Model %s price not configured. Go to System Settings → Operation Settings to enable self-use mode, or configure the model price in System Settings → Group & Model Pricing.",
modelName, modelName,
)
key = i18n.MsgRelayModelPriceNotConfiguredAdmin
}
return fmt.Errorf(
"模型 %s 的价格尚未由管理员配置,暂时无法使用,请联系站点管理员开启该模型;"+
"Model %s has not been priced by the administrator yet. Please contact the site administrator to enable this model.",
modelName, modelName,
)
return errors.New(i18n.T(c, key, map[string]any{"Model": modelName}))
}
// https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration
const claudeCacheCreation1hMultiplier = 6 / 3.75
// defaultTieredPreConsumeMaxTokens is the fallback completion-token estimate
// used for tiered expression pre-consume when the client omits max_tokens, so
// the pre-consumed quota still reflects a plausible output cost in paid groups.
const defaultTieredPreConsumeMaxTokens = 8192
// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) hosttypes.GroupRatioInfo {
if autoGroup, exists := ctx.Get("auto_group"); exists {
@@ -54,6 +44,16 @@ func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) hostty
return model.EffectiveGroupRatio(ctx, relayInfo.UserGroup, relayInfo.UsingGroup)
}
func deepSeekOfficialGroupRatio(info *relaycommon.RelayInfo, ratio hosttypes.GroupRatioInfo) hosttypes.GroupRatioInfo {
if info == nil || info.UsingGroup != billing_setting.DeepSeekOfficialGroup {
return ratio
}
return hosttypes.GroupRatioInfo{
GroupRatio: billing_setting.DeepSeekOfficialGroupRatio,
GroupSpecialRatio: -1,
}
}
func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (hosttypes.PriceData, error) {
if info != nil {
if matched := resolveBillingModelName(info.GetOriginModelName()); matched != "" && matched != info.OriginModelName {
@@ -65,10 +65,16 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
groupRatioInfo := HandleGroupRatio(c, info)
if info.UsingGroup == billing_setting.DeepSeekOfficialGroup {
if _, ok := billing_setting.GetGroupBillingExpr(info.UsingGroup, billingModelName); ok {
groupRatioInfo = deepSeekOfficialGroupRatio(info, groupRatioInfo)
}
}
// Check if this model uses tiered_expr billing
_, hasGroupExpr := billing_setting.GetGroupBillingExpr(info.UsingGroup, billingModelName)
if hasGroupExpr || billing_setting.GetBillingMode(billingModelName) == billing_setting.BillingModeTieredExpr {
return modelPriceHelperTiered(c, info, billingModelName, promptTokens, meta, groupRatioInfo)
return modelPriceHelperTiered(c, info, billingModelName, promptTokens, groupRatioInfo)
}
var preConsumedQuota int
@@ -83,10 +89,11 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
var audioCompletionRatio float64
var freeModel bool
if !usePrice {
preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota)
if meta.MaxTokens != 0 {
preConsumedTokens += meta.MaxTokens
preConsumeMultiplier, err := operation_setting.InputPreConsumeMultiplier()
if err != nil {
return hosttypes.PriceData{}, err
}
preConsumedTokens := float64(promptTokens) * preConsumeMultiplier
var success bool
var matchName string
modelRatio, success, matchName = ratio_setting.GetModelRatio(billingModelName)
@@ -96,7 +103,7 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
acceptUnsetRatio = true
}
if !acceptUnsetRatio {
return hosttypes.PriceData{}, modelPriceNotConfiguredError(matchName, info.UserId)
return hosttypes.PriceData{}, modelPriceNotConfiguredError(c, matchName, info.UserId)
}
}
completionRatio = ratio_setting.GetCompletionRatio(billingModelName)
@@ -109,13 +116,13 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
audioRatio = ratio_setting.GetAudioRatio(billingModelName)
audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(billingModelName)
ratio := modelRatio * groupRatioInfo.GroupRatio
quota, err := common.QuotaFromFloatStrict(float64(preConsumedTokens) * ratio)
quota, err := common.QuotaFromFloatStrict(preConsumedTokens * ratio)
if err != nil {
return hosttypes.PriceData{}, err
}
preConsumedQuota = quota
if _, image := info.Request.(*dto.ImageRequest); image {
info.ImageQuotaBeforeGroup = float64(preConsumedTokens) * modelRatio
info.ImageQuotaBeforeGroup = preConsumedTokens * modelRatio
}
} else {
if meta.ImagePriceRatio != 0 {
@@ -167,25 +174,13 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
}
}
if request, image := info.Request.(*dto.ImageRequest); image {
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
count, err := request.ImageCount(channelType == constant.ChannelTypeAli)
count, err := request.ImageCount(false)
if err != nil {
return hosttypes.PriceData{}, err
}
if usePrice || channelType == constant.ChannelTypeAli {
if usePrice {
priceData.AddOtherRatio("n", float64(count))
}
if channelType == constant.ChannelTypeAli && request.BillingParameters != nil && request.BillingParameters.PromptExtend != nil && *request.BillingParameters.PromptExtend {
// Resolve only routing identity; do not initialize ChannelMeta on the
// real request, which also distinguishes the first channel attempt.
mapped := &relaycommon.RelayInfo{OriginModelName: info.OriginModelName, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: info.OriginModelName}}
if err := ModelMappedHelper(c, mapped, nil); err != nil {
return hosttypes.PriceData{}, err
}
if strings.Contains(mapped.UpstreamModelName, "z-image") {
priceData.AddOtherRatio("prompt_extend", common.ZImagePromptExtendMultiplier)
}
}
if !usePrice {
quota, err := common.QuotaFromFloatStrict(priceData.ApplyOtherRatiosToFloat(info.ImageQuotaBeforeGroup * groupRatioInfo.GroupRatio))
if err != nil {
@@ -232,7 +227,7 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (hostt
acceptUnsetRatio = true
}
if !ratioSuccess && !acceptUnsetRatio {
return hosttypes.PriceData{}, modelPriceNotConfiguredError(matchName, info.UserId)
return hosttypes.PriceData{}, modelPriceNotConfiguredError(c, matchName, info.UserId)
}
}
}
@@ -314,6 +309,11 @@ func resolveBillingModelName(origin string) string {
}
candidates = append(candidates, hostreasoning.CanonicalBillingModelNames(origin)...)
base := hostreasoning.BaseModelName(origin)
// DeepSeek thinking tails (-none/-low/-max) bill at the base model's price,
// with or without @ modifiers after them.
if deepSeekBase, _, _, ok := reasoning.ParseDeepSeekV4ThinkingSuffix(base); ok && !model_setting.ShouldPreserveThinkingSuffix(base) {
candidates = append(candidates, deepSeekBase)
}
candidates = append(candidates, base)
seen := make(map[string]struct{}, len(candidates))
@@ -337,19 +337,23 @@ func resolveBillingModelName(origin string) string {
return matched
}
func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billingModelName string, promptTokens int, meta *types.TokenCountMeta, groupRatioInfo hosttypes.GroupRatioInfo) (hosttypes.PriceData, error) {
func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billingModelName string, promptTokens int, groupRatioInfo hosttypes.GroupRatioInfo) (hosttypes.PriceData, error) {
exprStr, ok := billing_setting.GetBillingExprForGroup(info.UsingGroup, billingModelName)
if !ok {
return hosttypes.PriceData{}, fmt.Errorf("model %s is configured as tiered_expr but has no billing expression", billingModelName)
}
return modelPriceHelperTieredWithExpr(c, info, billingModelName, promptTokens, groupRatioInfo, exprStr)
}
func modelPriceHelperTieredWithExpr(c *gin.Context, info *relaycommon.RelayInfo, billingModelName string, promptTokens int, groupRatioInfo hosttypes.GroupRatioInfo, exprStr string) (hosttypes.PriceData, error) {
exprHash := billingexpr.ExprHashString(exprStr)
if info.RelayFormat == types.RelayFormatOpenAIRealtime && billingexpr.UsesFixedPricingByHash(exprStr, exprHash) {
return hosttypes.PriceData{}, fmt.Errorf("fixed pricing is not supported for Realtime requests")
}
estimatedCompletionTokens := meta.MaxTokens
if estimatedCompletionTokens == 0 && groupRatioInfo.GroupRatio != 0 {
estimatedCompletionTokens = defaultTieredPreConsumeMaxTokens
preConsumeMultiplier, err := operation_setting.InputPreConsumeMultiplier()
if err != nil {
return hosttypes.PriceData{}, err
}
requestInput, err := ResolveIncomingBillingExprRequestInput(c, info)
@@ -365,7 +369,7 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
rawCost, trace, err := billingexpr.RunExprByHashWithRequest(exprStr, exprHash, billingexpr.TokenParams{
P: float64(promptTokens),
C: float64(estimatedCompletionTokens),
C: 0,
Len: float64(promptTokens),
}, requestInput)
if err != nil {
@@ -374,6 +378,11 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
// Expression coefficients are $/1M tokens prices; convert to quota the same way per-call billing does.
quotaBeforeGroup := rawCost / 1_000_000 * common.QuotaPerUnit
// Scale the reservation, preserving the expression's context-length tier
// and leaving actual settlement and fixed request prices unchanged.
if trace.BillingUnit != billingexpr.BillingUnitRequest {
quotaBeforeGroup *= preConsumeMultiplier
}
preConsumedQuota, err := billingexpr.QuotaRoundStrict(quotaBeforeGroup * groupRatioInfo.GroupRatio)
if err != nil {
return hosttypes.PriceData{}, err
@@ -395,7 +404,8 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
ExprHash: exprHash,
GroupRatio: groupRatioInfo.GroupRatio,
EstimatedPromptTokens: promptTokens,
EstimatedCompletionTokens: estimatedCompletionTokens,
EstimatedCompletionTokens: 0,
PreConsumeMultiplier: preConsumeMultiplier,
EstimatedQuotaBeforeGroup: quotaBeforeGroup,
EstimatedQuotaAfterGroup: preConsumedQuota,
EstimatedTier: trace.MatchedTier,
+172 -33
View File
@@ -1,15 +1,19 @@
package helper
import (
"fmt"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/pkg/billingexpr"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/config"
"github.com/QuantumNous/new-api/setting/model_setting"
@@ -20,6 +24,39 @@ import (
"github.com/stretchr/testify/require"
)
func TestDeepSeekOfficialGroupUsesFixed045Ratio(t *testing.T) {
savedRatios := ratio_setting.GroupRatio2JSONString()
t.Cleanup(func() { require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(savedRatios)) })
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"Deepseek官Key":0.3,"default":1}`))
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
OriginModelName: "deepseek-v4.1-flash",
UserGroup: "default",
UsingGroup: billing_setting.DeepSeekOfficialGroup,
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeOpenAI},
BillingRequestInput: &billingexpr.RequestInput{
Body: []byte(`{}`),
},
}
price, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
require.NoError(t, err)
assert.Equal(t, billing_setting.DeepSeekOfficialGroupRatio, price.GroupRatioInfo.GroupRatio)
assert.Equal(t, billing_setting.DeepSeekOfficialGroupRatio, info.TieredBillingSnapshot.GroupRatio)
assert.Equal(t, billing_setting.BillingModeTieredExpr, info.TieredBillingSnapshot.BillingMode)
expression, ok := billing_setting.GetGroupBillingExpr(info.UsingGroup, info.GetBillingModelName())
require.True(t, ok)
assert.Equal(t, expression, info.TieredBillingSnapshot.ExprString)
info.ChannelMeta.ChannelType = constant.ChannelTypeDeepSeek
_, err = ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
require.NoError(t, err)
assert.Equal(t, expression, info.TieredBillingSnapshot.ExprString)
_, hasDomesticOverride := billing_setting.GetGroupBillingExpr("国模", info.GetBillingModelName())
assert.False(t, hasDomesticOverride)
}
func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -109,7 +146,7 @@ func TestFixedPricePreConsumeAndRealtimeRejection(t *testing.T) {
}
}
func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
func TestModelPriceHelperTieredInputPreConsumeMultiplier(t *testing.T) {
gin.SetMode(gin.TestMode)
saved := map[string]string{}
@@ -123,41 +160,25 @@ func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{"tiered-fallback-model":"tiered_expr"}`,
"billing_setting.billing_expr": `{"tiered-fallback-model":"tier(\"base\", p * 3 + c * 15)"}`,
"billing_setting.billing_expr": `{"tiered-fallback-model":"len <= 1200 ? tier(\"base\", p * 3 + c * 15) : tier(\"long\", p * 30 + c * 150)"}`,
"group_ratio_setting.group_ratio": `{"default":1,"free":0}`,
}))
const promptTokens = 1000
cases := []struct {
name string
group string
maxTokens int
expected int
name string
group string
prompt int
maxTokens int
multiplier float64
expected int
}{
{
// max_tokens omitted in a paid group -> fall back to 8192 completion tokens.
// p*3 + c*15 = 1000*3 + 8192*15 = 125880 -> /1e6 * 500000 = 62940
name: "non-free group falls back to 8192 completion tokens",
group: "default",
maxTokens: 0,
expected: 62940,
},
{
// explicit max_tokens is used verbatim, no fallback.
// 1000*3 + 100*15 = 4500 -> /1e6 * 500000 = 2250
name: "explicit max_tokens is used verbatim",
group: "default",
maxTokens: 100,
expected: 2250,
},
{
// free group (ratio 0) stays zero; fallback is gated on non-zero group ratio.
name: "free group stays zero without fallback",
group: "free",
maxTokens: 0,
expected: 0,
},
{"default uses only input", "default", 1000, 0, 1, 1500},
{"explicit output limit does not increase reservation", "default", 1000, 100, 1, 1500},
{"fraction below one", "default", 1000, 0, 0.5, 750},
{"fraction above one preserves context tier", "default", 1000, 0, 2.5, 3750},
{"small input has no token floor", "default", 100, 0, 1, 150},
{"zero input", "default", 0, 100, 1, 0},
{"free group", "free", 1000, 0, 2.5, 0},
}
for _, tc := range cases {
@@ -180,9 +201,18 @@ func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
},
}
priceData, err := ModelPriceHelper(ctx, info, promptTokens, &types.TokenCountMeta{MaxTokens: tc.maxTokens})
operation_setting.GetQuotaSetting().PreConsumeMultiplier = tc.multiplier
priceData, err := ModelPriceHelper(ctx, info, tc.prompt, &types.TokenCountMeta{MaxTokens: tc.maxTokens})
require.NoError(t, err)
require.Equal(t, tc.expected, priceData.QuotaToPreConsume)
assert.Equal(t, tc.expected, priceData.QuotaToPreConsume)
require.NotNil(t, info.TieredBillingSnapshot)
actual, err := billingexpr.ComputeTieredQuotaWithRequest(info.TieredBillingSnapshot, billingexpr.TokenParams{P: 1000, C: 100, Len: 1000}, *info.BillingRequestInput)
require.NoError(t, err)
wantActual := 2250
if tc.group == "free" {
wantActual = 0
}
assert.Equal(t, wantActual, actual.ActualQuotaAfterGroup, "reservation multiplier must not change settlement")
})
}
}
@@ -650,6 +680,60 @@ func TestModelPriceHelperPreservesGpt51CodexMaxIdentity(t *testing.T) {
assert.Equal(t, 1.75, priceData.ModelRatio)
}
// DeepSeek thinking tails and @ modifiers bill at the base model's built-in
// expression; a name on the thinking-suffix blacklist stays opaque and needs
// its own price.
func TestModelPriceHelperDeepSeekThinkingTailsBillAtBase(t *testing.T) {
gin.SetMode(gin.TestMode)
require.NoError(t, i18n.Init())
saved := map[string]string{}
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
saved[key] = value
return nil
}))
t.Cleanup(func() { require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) })
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{}`,
"billing_setting.billing_expr": `{}`,
}))
oldSelfUse := operation_setting.SelfUseModeEnabled
operation_setting.SelfUseModeEnabled = false
t.Cleanup(func() { operation_setting.SelfUseModeEnabled = oldSelfUse })
settings := model_setting.GetGlobalSettings()
originalBlacklist := append([]string(nil), settings.ThinkingModelBlacklist...)
t.Cleanup(func() { settings.ThinkingModelBlacklist = originalBlacklist })
settings.ThinkingModelBlacklist = append(originalBlacklist, "deepseek-v4-pro-max")
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
ctx.Set("group", "default")
for origin, base := range map[string]string{
"deepseek-flash-none": "deepseek-flash",
"deepseek-flash-low": "deepseek-flash",
"deepseek-flash-max": "deepseek-flash",
"deepseek-v4-pro-low": "deepseek-v4-pro",
"deepseek-flash@effort:max": "deepseek-flash",
"deepseek-flash-max@effort:low": "deepseek-flash",
"deepseek-v4-pro@thinking:off": "deepseek-v4-pro",
} {
info := &relaycommon.RelayInfo{
OriginModelName: origin, UserGroup: "default", UsingGroup: "default",
BillingRequestInput: &billingexpr.RequestInput{Body: []byte(`{}`)},
}
_, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
require.NoError(t, err, origin)
assert.Equal(t, base, info.BillingModelName, origin)
expression, _ := billing_setting.GetBuiltinBillingExpr(base)
require.NotNil(t, info.TieredBillingSnapshot, origin)
assert.Equal(t, expression, info.TieredBillingSnapshot.ExprString, origin)
}
info := &relaycommon.RelayInfo{OriginModelName: "deepseek-v4-pro-max", UserGroup: "default", UsingGroup: "default"}
_, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
require.Error(t, err)
assert.Equal(t, "deepseek-v4-pro-max", info.GetBillingModelName())
}
func TestModelPriceHelperNativeGeminiNoThinkingDoesNotAliasBillingModel(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -695,3 +779,58 @@ func TestModelPriceHelperNativeGeminiNoThinkingDoesNotAliasBillingModel(t *testi
assert.Equal(t, 1.25, priceData.ModelRatio)
assert.NotEqual(t, 37.5, priceData.ModelRatio)
}
func TestInputPreConsumeMultiplierLegacyAndRequestPrices(t *testing.T) {
previous := config.GlobalConfig.ExportAllConfigs()
previousRatios, previousPrices := ratio_setting.ModelRatio2JSONString(), ratio_setting.ModelPrice2JSONString()
t.Cleanup(func() {
require.NoError(t, config.GlobalConfig.LoadFromDB(previous))
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(previousRatios))
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(previousPrices))
})
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"legacy-input-policy":1.5}`))
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"request-input-policy":0.01}`))
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{"fixed-input-policy":"tiered_expr","image-input-policy":"tiered_expr"}`,
"billing_setting.billing_expr": `{"fixed-input-policy":"tier(\"request\", fixed(0.01))","image-input-policy":"tier(\"image\", p * 3) * image_count"}`,
"group_ratio_setting.group_ratio": `{"default":1}`,
}))
for _, tc := range []struct {
model string
multiplier float64
prompt, want int
}{
{"legacy-input-policy", 0.5, 100, 75},
{"legacy-input-policy", 2.5, 100, 375},
{"legacy-input-policy", 1, 0, 0},
{"request-input-policy", 2.5, 100, 5000},
{"fixed-input-policy", 2.5, 100, 5000},
{"image-input-policy", 2.5, 100, 375},
} {
t.Run(tc.model+"/"+fmt.Sprint(tc.multiplier, "/", tc.prompt), func(t *testing.T) {
operation_setting.GetQuotaSetting().PreConsumeMultiplier = tc.multiplier
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{OriginModelName: tc.model, UserGroup: "default", UsingGroup: "default", BillingRequestInput: &billingexpr.RequestInput{}}
price, err := ModelPriceHelper(ctx, info, tc.prompt, &types.TokenCountMeta{MaxTokens: 10000})
require.NoError(t, err)
assert.Equal(t, tc.want, price.QuotaToPreConsume)
if tc.model == "image-input-policy" {
reservation := &priceTestReservation{held: price.QuotaToPreConsume}
info.Billing = reservation
operation_setting.GetQuotaSetting().PreConsumeMultiplier = 10
require.Nil(t, service.PrepareImageBillingForRequest(ctx, info, 2))
assert.Equal(t, 750, reservation.held, "retry must reuse the original fractional multiplier")
}
})
}
}
// priceTestReservation observes the reservation requested before image submission.
type priceTestReservation struct{ held int }
func (s *priceTestReservation) Settle(int) error { return nil }
func (s *priceTestReservation) Refund(*gin.Context) {}
func (s *priceTestReservation) NeedsRefund() bool { return false }
func (s *priceTestReservation) GetPreConsumedQuota() int { return s.held }
func (s *priceTestReservation) Reserve(quota int) error { s.held = max(s.held, quota); return nil }
+21
View File
@@ -161,6 +161,27 @@ func TestApplyReasoningModelSuffixGeminiNoThinkingWhenAdapterEnabled(t *testing.
assert.Equal(t, "none", info.ReasoningConversion.Effort)
}
func TestApplyReasoningModelSuffixClaudeEffortTailKeptWhenAdapterDisabled(t *testing.T) {
settings := model_setting.GetClaudeSettings()
original := settings.ThinkingAdapterEnabled
t.Cleanup(func() { settings.ThinkingAdapterEnabled = original })
settings.ThinkingAdapterEnabled = false
req := &dto.ClaudeRequest{Model: "claude-opus-5-5-high"}
info := &relaycommon.RelayInfo{
OriginModelName: "claude-opus-5-5-high",
Request: req,
ChannelMeta: &relaycommon.ChannelMeta{
UpstreamModelName: "claude-opus-5-5-high",
},
}
mustApplyReasoningModelSuffix(t, info, req)
assert.Equal(t, "claude-opus-5-5-high", info.UpstreamModelName)
assert.Equal(t, "claude-opus-5-5-high", req.Model)
assert.Nil(t, info.ReasoningConversion)
}
func TestApplyReasoningModelSuffixPreservesEffortTailModelID(t *testing.T) {
info := &relaycommon.RelayInfo{
OriginModelName: "qwen-max",
+9 -8
View File
@@ -161,7 +161,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
gopool.Go(func() {
defer func() {
if r := recover(); r != nil {
logger.LogError(c, fmt.Sprintf("ping goroutine panic: %v", r))
logger.LogError(c, common.LogText("ping goroutine panic: %v", r))
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("ping panic: %v", r))
stop()
}
@@ -185,7 +185,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
err = PingData(c)
}()
if err != nil {
logger.LogError(c, "ping data error: "+err.Error())
logger.LogError(c, common.LogText("ping data error: %s", err.Error()))
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPingFail, err)
return
}
@@ -198,7 +198,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
// 监听客户端断开连接
return
case <-pingTimeout.C:
logger.LogError(c, "ping goroutine max duration reached")
logger.LogError(c, common.LogText("ping goroutine max duration reached"))
return
}
}
@@ -211,7 +211,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
gopool.Go(func() {
defer func() {
if r := recover(); r != nil {
logger.LogError(c, fmt.Sprintf("data handler goroutine panic: %v", r))
logger.LogError(c, common.LogText("data handler goroutine panic: %v", r))
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("handler panic: %v", r))
}
stop()
@@ -238,7 +238,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
defer func() {
close(dataChan)
if r := recover(); r != nil {
logger.LogError(c, fmt.Sprintf("scanner goroutine panic: %v", r))
logger.LogError(c, common.LogText("scanner goroutine panic: %v", r))
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("scanner panic: %v", r))
}
stop()
@@ -272,6 +272,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
continue
}
if !strings.HasPrefix(data, "[DONE]") {
info.UpstreamTiming.Mark("first_sse_data")
info.SetFirstResponseTime()
info.ReceivedResponseCount++
@@ -291,7 +292,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
if err := scanner.Err(); err != nil {
if err != io.EOF {
logger.LogError(c, "scanner error: "+err.Error())
logger.LogError(c, common.LogText("scanner error: %s", err.Error()))
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonScannerErr, err)
}
}
@@ -337,8 +338,8 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
cleanup(true)
if info.StreamStatus.IsNormalEnd() && !info.StreamStatus.HasErrors() {
logger.LogInfo(c, fmt.Sprintf("stream ended: %s", info.StreamStatus.Summary()))
logger.LogInfo(c, common.LogText("stream ended: %s", info.StreamStatus.Summary()))
} else {
logger.LogError(c, fmt.Sprintf("stream ended: %s, received=%d", info.StreamStatus.Summary(), info.ReceivedResponseCount))
logger.LogError(c, common.LogText("stream ended: %s, received=%d", info.StreamStatus.Summary(), info.ReceivedResponseCount))
}
}
+27 -4
View File
@@ -10,7 +10,6 @@ import (
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
@@ -42,6 +41,8 @@ func GetAndValidateRequest(c *gin.Context, format types.RelayFormat) (request dt
request, err = GetAndValidateResponsesCompactionRequest(c)
case types.RelayFormatOpenAIAlphaSearch:
request, err = GetAndValidateAlphaSearchRequest(c)
case types.RelayFormatDeepSeekTools:
request, err = GetAndValidateDeepSeekToolsRequest(c)
case types.RelayFormatOpenAIImage:
request, err = GetAndValidOpenAIImageRequest(c, relayMode)
@@ -85,7 +86,7 @@ func GetAndValidateRerankRequest(c *gin.Context) (*dto.RerankRequest, error) {
var rerankRequest *dto.RerankRequest
err := common.UnmarshalBodyReusable(c, &rerankRequest)
if err != nil {
logger.LogError(c, fmt.Sprintf("getAndValidateTextRequest failed: %s", err.Error()))
logger.LogError(c, common.LogText("getAndValidateTextRequest failed: %s", err.Error()))
return nil, types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
@@ -102,7 +103,7 @@ func GetAndValidateEmbeddingRequest(c *gin.Context, relayMode int) (*dto.Embeddi
var embeddingRequest *dto.EmbeddingRequest
err := common.UnmarshalBodyReusable(c, &embeddingRequest)
if err != nil {
logger.LogError(c, fmt.Sprintf("getAndValidateTextRequest failed: %s", err.Error()))
logger.LogError(c, common.LogText("getAndValidateTextRequest failed: %s", err.Error()))
return nil, types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
}
@@ -171,6 +172,28 @@ func GetAndValidateAlphaSearchRequest(c *gin.Context) (*dto.AlphaSearchRequest,
return request, nil
}
// GetAndValidateDeepSeekToolsRequest keeps the raw body for forwarding. The
// upstream validates query/url itself, and no field changes the per-call price.
func GetAndValidateDeepSeekToolsRequest(c *gin.Context) (*dto.DeepSeekToolsRequest, error) {
request := &dto.DeepSeekToolsRequest{}
if err := common.UnmarshalBodyReusable(c, request); err != nil {
return nil, err
}
storage, err := common.GetBodyStorage(c)
if err != nil {
return nil, err
}
rawBody, err := storage.Bytes()
if err != nil {
return nil, err
}
if len(rawBody) == 0 {
return nil, errors.New("request body is empty")
}
request.RawBody = rawBody
return request, nil
}
func GetAndValidateResponsesCompactionRequest(c *gin.Context) (*dto.OpenAIResponsesCompactionRequest, error) {
request := &dto.OpenAIResponsesCompactionRequest{}
if err := common.UnmarshalBodyReusable(c, request); err != nil {
@@ -283,7 +306,7 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
}
imageRequest.BillingParameters = parameters
}
if _, err := imageRequest.ImageCount(common.GetContextKeyInt(c, constant.ContextKeyChannelType) == constant.ChannelTypeAli); err != nil {
if _, err := imageRequest.ImageCount(false); err != nil {
return nil, err
}
return imageRequest, nil