52 lines
2.1 KiB
Go
52 lines
2.1 KiB
Go
package helper
|
|
|
|
import (
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/QuantumNous/new-api/setting/config"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestGroupModelPricingPreservesOtherGroupsAndSnapshot(t *testing.T) {
|
|
saved := map[string]string{}
|
|
require.NoError(t, config.GlobalConfig.SaveToDB(func(k, v string) error { saved[k] = v; 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": `{"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}`,
|
|
}))
|
|
for _, tc := range []struct {
|
|
group string
|
|
expected int
|
|
cost float64
|
|
}{
|
|
{"example-premium", 1000, 8500},
|
|
{"example-standard", 1500, 6400},
|
|
} {
|
|
t.Run(tc.group, func(t *testing.T) {
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
ctx.Set("auto_group", tc.group)
|
|
info := &relaycommon.RelayInfo{
|
|
OriginModelName: "gpt-5.6-sol", UserGroup: "default", UsingGroup: "auto",
|
|
BillingRequestInput: &billingexpr.RequestInput{Body: []byte(`{}`)},
|
|
}
|
|
data, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{MaxTokens: 100})
|
|
require.NoError(t, err)
|
|
require.Equal(t, tc.expected, data.QuotaToPreConsume)
|
|
require.Equal(t, tc.group, info.UsingGroup)
|
|
raw, _, err := billingexpr.RunExprWithRequest(info.TieredBillingSnapshot.ExprString,
|
|
billingexpr.TokenParams{P: 1000, C: 100, CR: 1000, Len: 2000}, billingexpr.RequestInput{})
|
|
require.NoError(t, err)
|
|
require.Equal(t, tc.cost, raw)
|
|
})
|
|
}
|
|
}
|