112 lines
4.7 KiB
Go
112 lines
4.7 KiB
Go
package billing_setting_test
|
|
|
|
import (
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/controller"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"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/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestGPT6AstraBuiltinBilling(t *testing.T) {
|
|
settings := config.GlobalConfig.Get("billing_setting").(*billing_setting.BillingSetting)
|
|
saved := *settings
|
|
savedRatios, savedPrices := ratio_setting.ModelRatio2JSONString(), ratio_setting.ModelPrice2JSONString()
|
|
savedOptions := common.OptionMap
|
|
t.Cleanup(func() {
|
|
*settings, common.OptionMap = saved, savedOptions
|
|
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedRatios))
|
|
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(savedPrices))
|
|
})
|
|
common.OptionMap = map[string]string{"billing_setting.billing_mode": `{}`, "billing_setting.billing_expr": `{}`}
|
|
require.NoError(t, config.GlobalConfig.LoadFromDB(common.OptionMap))
|
|
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`))
|
|
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
|
|
assert.Equal(t, billing_setting.BillingModeTieredExpr, billing_setting.GetBillingMode("gpt-6-astra"))
|
|
expression, ok := billing_setting.GetBillingExpr("gpt-6-astra")
|
|
require.True(t, ok)
|
|
|
|
for _, tc := range []struct {
|
|
name string
|
|
input, output, cached, written int
|
|
request string
|
|
quota int
|
|
}{
|
|
{"standard", 1000, 100, 0, 0, `{}`, 7500},
|
|
{"client flex cannot discount standard pricing", 1000, 100, 0, 0, `{"service_tier":"flex"}`, 7500},
|
|
{"cache at context boundary", 272000, 1000, 200000, 20000, `{}`, 510000},
|
|
{"whole request above boundary", 272001, 1000, 200000, 20000, `{}`, 1007510},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
usage := &dto.Usage{
|
|
PromptTokens: tc.input, CompletionTokens: tc.output,
|
|
PromptTokensDetails: dto.InputTokenDetails{CachedTokens: tc.cached, CacheWriteTokens: tc.written},
|
|
}
|
|
params := service.BuildTieredTokenParams(usage, false, billingexpr.UsedVars(expression))
|
|
result, err := billingexpr.ComputeTieredQuotaWithRequest(&billingexpr.BillingSnapshot{
|
|
ExprString: expression, GroupRatio: 1, QuotaPerUnit: 500000,
|
|
}, params, billingexpr.RequestInput{Body: []byte(tc.request)})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, tc.quota, result.ActualQuotaAfterGroup)
|
|
})
|
|
}
|
|
|
|
t.Run("admin options expose defaults without persisting them", func(t *testing.T) {
|
|
recorder := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(recorder)
|
|
controller.GetOptions(ctx)
|
|
var response struct {
|
|
Success bool
|
|
Data []model.Option
|
|
}
|
|
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
|
require.True(t, response.Success)
|
|
found := map[string]string{}
|
|
for _, option := range response.Data {
|
|
if _, ok := common.OptionMap[option.Key]; ok {
|
|
assert.NotContains(t, found, option.Key)
|
|
var values map[string]string
|
|
require.NoError(t, common.UnmarshalJsonStr(option.Value, &values))
|
|
found[option.Key] = values["gpt-6-astra"]
|
|
assert.Equal(t, `{}`, common.OptionMap[option.Key])
|
|
}
|
|
}
|
|
assert.Equal(t, map[string]string{"billing_setting.billing_mode": "tiered_expr", "billing_setting.billing_expr": expression}, found)
|
|
})
|
|
|
|
for _, tc := range []struct {
|
|
name, mode, expr, ratios, prices, wantMode string
|
|
}{
|
|
{"custom expression overrides legacy price", "tiered_expr", "p * 7", `{"gpt-6-astra":8}`, `{}`, "tiered_expr"},
|
|
{"explicit ratio mode", "ratio", "", `{}`, `{}`, "ratio"},
|
|
{"existing free token price", "", "", `{"gpt-6-astra":0}`, `{}`, "ratio"},
|
|
{"existing per-call price", "", "", `{}`, `{"gpt-6-astra":0.1}`, "ratio"},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
*settings = billing_setting.BillingSetting{BillingMode: map[string]string{}, BillingExpr: map[string]string{}}
|
|
if tc.mode != "" {
|
|
settings.BillingMode["gpt-6-astra"] = tc.mode
|
|
}
|
|
if tc.expr != "" {
|
|
settings.BillingExpr["gpt-6-astra"] = tc.expr
|
|
}
|
|
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(tc.ratios))
|
|
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(tc.prices))
|
|
assert.Equal(t, tc.wantMode, billing_setting.GetBillingMode("gpt-6-astra"))
|
|
actual, ok := billing_setting.GetBillingExpr("gpt-6-astra")
|
|
assert.Equal(t, tc.expr, actual)
|
|
assert.Equal(t, tc.expr != "", ok)
|
|
})
|
|
}
|
|
}
|