feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -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
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
@@ -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 }
|
||||
@@ -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",
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user