feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -1,15 +1,25 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -177,3 +187,139 @@ export function parseTaskResult() { return {}; }
|
||||
assert.Contains(t, unresolvable.Body.String(), `"success":false`)
|
||||
assert.Contains(t, unresolvable.Body.String(), "no task plugin usage schema")
|
||||
}
|
||||
|
||||
// Covers configuration persistence, validation, reservation and reconciliation
|
||||
// through the same public entry points used by the relay.
|
||||
func TestPreConsumePolicyDatabaseMatrix(t *testing.T) {
|
||||
previousConfig := config.GlobalConfig.ExportAllConfigs()
|
||||
previousUnit, previousBatch := common.QuotaPerUnit, common.BatchUpdateEnabled
|
||||
common.QuotaPerUnit, common.BatchUpdateEnabled = 500000, false
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(previousConfig))
|
||||
common.QuotaPerUnit, common.BatchUpdateEnabled = previousUnit, previousBatch
|
||||
})
|
||||
for _, dialect := range []struct{ kind, env string }{{"sqlite", ""}, {"mysql", "TEST_MYSQL_DSN"}, {"postgres", "TEST_POSTGRES_DSN"}} {
|
||||
t.Run(dialect.kind, func(t *testing.T) {
|
||||
if dialect.env != "" && os.Getenv(dialect.env) == "" {
|
||||
t.Skip("set " + dialect.env + " to run this database")
|
||||
}
|
||||
db := modelManagementDB(t, dialect.kind, os.Getenv(dialect.env))
|
||||
require.NoError(t, db.AutoMigrate(&model.Token{}))
|
||||
for key, value := range map[string]string{
|
||||
"quota_setting.trust_quota_usd": "10.5",
|
||||
"quota_setting.pre_consume_multiplier": "1.5",
|
||||
} {
|
||||
var response struct{ Success bool }
|
||||
modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: key, Value: value}, &response)
|
||||
require.True(t, response.Success)
|
||||
var saved model.Option
|
||||
require.NoError(t, db.Where(&model.Option{Key: key}).First(&saved).Error)
|
||||
assert.Equal(t, value, saved.Value)
|
||||
}
|
||||
var savedOptions []model.Option
|
||||
require.NoError(t, db.Find(&savedOptions).Error)
|
||||
stored := map[string]string{}
|
||||
for _, option := range savedOptions {
|
||||
stored[option.Key] = option.Value
|
||||
}
|
||||
for range 2 {
|
||||
operation_setting.GetQuotaSetting().TrustQuotaUSD = 10
|
||||
operation_setting.GetQuotaSetting().PreConsumeMultiplier = 1
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(stored))
|
||||
assert.Equal(t, 10.5, operation_setting.GetQuotaSetting().TrustQuotaUSD)
|
||||
assert.Equal(t, 1.5, operation_setting.GetQuotaSetting().PreConsumeMultiplier)
|
||||
}
|
||||
for _, tc := range []struct{ key, value string }{
|
||||
{"trust_quota_usd", "-1"}, {"trust_quota_usd", "NaN"}, {"trust_quota_usd", "+Inf"},
|
||||
{"pre_consume_multiplier", "0"}, {"pre_consume_multiplier", "-0.5"},
|
||||
{"pre_consume_multiplier", "NaN"}, {"pre_consume_multiplier", "Inf"},
|
||||
{"pre_consume_multiplier", "1e309"}, {"pre_consume_multiplier", ""},
|
||||
} {
|
||||
key := "quota_setting." + tc.key
|
||||
var response struct{ Success bool }
|
||||
modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: key, Value: tc.value}, &response)
|
||||
assert.False(t, response.Success, "%s=%q", key, tc.value)
|
||||
var saved model.Option
|
||||
require.NoError(t, db.Where(&model.Option{Key: key}).First(&saved).Error)
|
||||
assert.Equal(t, stored[key], saved.Value, "invalid saves must preserve the previous value")
|
||||
}
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
"billing_setting.billing_mode": `{"policy-model":"tiered_expr"}`,
|
||||
"billing_setting.billing_expr": `{"policy-model":"tier(\"base\", p * 3 + c * 15)"}`,
|
||||
"group_ratio_setting.group_ratio": `{"default":1}`,
|
||||
}))
|
||||
cases := []struct {
|
||||
name string
|
||||
threshold, multiplier float64
|
||||
wallet, token int
|
||||
force, unlimited bool
|
||||
wantHeld int
|
||||
}{
|
||||
{"above threshold", 10, 1, 5500000, 5500000, false, false, 0},
|
||||
{"custom threshold", 20, 1, 5500000, 5500000, false, false, 1500},
|
||||
{"zero disables bypass", 0, 1, 5500000, 5500000, false, false, 1500},
|
||||
{"wallet equals threshold", 10, 1, 5000000, 5500000, false, false, 1500},
|
||||
{"token equals threshold", 10, 1, 5500000, 5000000, false, false, 1500},
|
||||
{"fractional threshold", 10.5, 1, 5250001, 5250001, false, false, 0},
|
||||
{"unlimited token", 10, 1, 5500000, 0, false, true, 0},
|
||||
{"forced reservation", 10, 1, 5500000, 5500000, true, false, 1500},
|
||||
{"half input cost", 0, 0.5, 100000, 100000, false, false, 750},
|
||||
{"fractional multiple refunds excess", 0, 2.5, 100000, 100000, false, false, 3750},
|
||||
}
|
||||
for i, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
require.NoError(t, model.UpdateOption("quota_setting.trust_quota_usd", strconv.FormatFloat(tc.threshold, 'f', -1, 64)))
|
||||
require.NoError(t, model.UpdateOption("quota_setting.pre_consume_multiplier", strconv.FormatFloat(tc.multiplier, 'f', -1, 64)))
|
||||
user := model.User{Username: fmt.Sprintf("policy-user-%d", i), Quota: tc.wallet, Group: "default", AffCode: fmt.Sprintf("policy-aff-%d", i)}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
token := model.Token{UserId: user.Id, Key: fmt.Sprintf("policy-token-%d", i), RemainQuota: tc.token, UnlimitedQuota: tc.unlimited}
|
||||
require.NoError(t, db.Create(&token).Error)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
ctx.Set("token_quota", tc.token)
|
||||
info := &relaycommon.RelayInfo{UserId: user.Id, TokenId: token.Id, TokenKey: token.Key, TokenUnlimited: tc.unlimited, ForcePreConsume: tc.force, OriginModelName: "policy-model", UserGroup: "default", UsingGroup: "default", BillingRequestInput: &billingexpr.RequestInput{}}
|
||||
info.UserSetting.BillingPreference = "wallet_only"
|
||||
price, err := helper.ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{MaxTokens: 10000})
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, service.PreConsumeBilling(ctx, price.QuotaToPreConsume, info))
|
||||
assert.Equal(t, tc.wantHeld, info.FinalPreConsumedQuota)
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
assert.Equal(t, tc.wallet-tc.wantHeld, user.Quota)
|
||||
if !tc.unlimited {
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
assert.Equal(t, tc.token-tc.wantHeld, token.RemainQuota)
|
||||
}
|
||||
_, actual, _ := service.TryTieredSettle(info, billingexpr.TokenParams{P: 1000, C: 100, Len: 1000})
|
||||
assert.Equal(t, 2250, actual)
|
||||
require.NoError(t, info.Billing.Settle(actual))
|
||||
require.NoError(t, info.Billing.Settle(actual))
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
assert.Equal(t, tc.wallet-actual, user.Quota)
|
||||
if !tc.unlimited {
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
assert.Equal(t, tc.token-actual, token.RemainQuota)
|
||||
}
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreConsumeMultiplierRejectsInvalidRuntimeAndOverflow(t *testing.T) {
|
||||
previous := config.GlobalConfig.ExportAllConfigs()
|
||||
t.Cleanup(func() { require.NoError(t, config.GlobalConfig.LoadFromDB(previous)) })
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
"billing_setting.billing_mode": `{"policy-overflow":"tiered_expr"}`,
|
||||
"billing_setting.billing_expr": `{"policy-overflow":"tier(\"base\", p * 3)"}`,
|
||||
"group_ratio_setting.group_ratio": `{"default":1}`,
|
||||
}))
|
||||
for _, multiplier := range []float64{0, -0.5, math.NaN(), math.Inf(1), math.MaxFloat64} {
|
||||
operation_setting.GetQuotaSetting().PreConsumeMultiplier = multiplier
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "policy-overflow", UserGroup: "default", UsingGroup: "default", BillingRequestInput: &billingexpr.RequestInput{}}
|
||||
_, err := helper.ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
|
||||
require.Error(t, err)
|
||||
assert.Nil(t, info.Billing)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user