1783 lines
69 KiB
Go
1783 lines
69 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"math"
|
|
"net/http/httptest"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/relaykit/relayconvert"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/QuantumNous/new-api/setting/operation_setting"
|
|
hosttypes "github.com/QuantumNous/new-api/types"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/shopspring/decimal"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// The configured DSNs must point at isolated test databases. Each dialect runs
|
|
// the real reservation, settlement and log paths with the same billing cases.
|
|
func TestFixedPriceBillingDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []struct {
|
|
name common.DatabaseType
|
|
env string
|
|
logEnv string
|
|
}{
|
|
{common.DatabaseTypeSQLite, "", ""},
|
|
{common.DatabaseTypeMySQL, "TEST_FIXED_MYSQL_DSN", "TEST_FIXED_MYSQL_LOG_DSN"},
|
|
{common.DatabaseTypePostgreSQL, "TEST_FIXED_POSTGRES_DSN", "TEST_FIXED_POSTGRES_LOG_DSN"},
|
|
} {
|
|
t.Run(string(dialect.name), func(t *testing.T) {
|
|
var driver gorm.Dialector = sqlite.Open(":memory:")
|
|
if dialect.env != "" {
|
|
dsn := os.Getenv(dialect.env)
|
|
if dsn == "" {
|
|
t.Skip(dialect.env + " is not configured")
|
|
}
|
|
if dialect.name == common.DatabaseTypeMySQL {
|
|
driver = mysql.Open(dsn)
|
|
} else {
|
|
driver = postgres.New(postgres.Config{DSN: dsn, PreferSimpleProtocol: true})
|
|
}
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{})
|
|
require.NoError(t, err)
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
sqlDB.SetMaxOpenConns(1)
|
|
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
|
|
logDB := db
|
|
if logDSN := os.Getenv(dialect.logEnv); logDSN != "" || dialect.name == common.DatabaseTypeSQLite {
|
|
var logDriver gorm.Dialector = sqlite.Open(":memory:")
|
|
if dialect.name == common.DatabaseTypeMySQL {
|
|
logDriver = mysql.Open(logDSN)
|
|
} else if dialect.name == common.DatabaseTypePostgreSQL {
|
|
logDriver = postgres.New(postgres.Config{DSN: logDSN, PreferSimpleProtocol: true})
|
|
}
|
|
logDB, err = gorm.Open(logDriver, &gorm.Config{})
|
|
require.NoError(t, err)
|
|
logSQL, err := logDB.DB()
|
|
require.NoError(t, err)
|
|
logSQL.SetMaxOpenConns(1)
|
|
t.Cleanup(func() { require.NoError(t, logSQL.Close()) })
|
|
}
|
|
oldDB, oldLogDB := model.DB, model.LOG_DB
|
|
oldMainType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
|
|
model.DB, model.LOG_DB = db, logDB
|
|
common.SetDatabaseTypes(dialect.name, dialect.name)
|
|
t.Cleanup(func() { model.DB, model.LOG_DB = oldDB, oldLogDB; common.SetDatabaseTypes(oldMainType, oldLogType) })
|
|
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}))
|
|
require.NoError(t, logDB.AutoMigrate(&model.Log{}))
|
|
versionQuery := "select version()"
|
|
if dialect.name == common.DatabaseTypeSQLite {
|
|
versionQuery = "select sqlite_version()"
|
|
}
|
|
var version string
|
|
require.NoError(t, db.Raw(versionQuery).Scan(&version).Error)
|
|
t.Logf("database: %s", version)
|
|
runFixedPriceAccountingCases(t, db, logDB)
|
|
t.Run("claude cache read compatibility", func(t *testing.T) {
|
|
runClaudeCacheReadAccountingCases(t, db, logDB)
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
func claudeCacheReadTestInfo() *relaycommon.RelayInfo {
|
|
return &relaycommon.RelayInfo{
|
|
OriginModelName: "deepseek-v4.1-flash", RelayFormat: types.RelayFormatClaude,
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelId: 22, UpstreamModelName: "deepseek-v4.1-flash",
|
|
ChannelOtherSettings: dto.ChannelOtherSettings{ClaudeInputIncludesCacheReadModels: []string{"deepseek-v4.1-flash"}},
|
|
},
|
|
PriceData: hosttypes.PriceData{ModelRatio: 1, CompletionRatio: 4, CacheRatio: 0.02, CacheCreationRatio: 1.25,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 0.162}},
|
|
StartTime: time.Now(),
|
|
}
|
|
}
|
|
|
|
func TestClaudeCacheReadCompatibility(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
prompt, cached, output int
|
|
wantQuota int
|
|
}{
|
|
{"screenshot 1", 275773, 275200, 287, 1170},
|
|
{"screenshot 2", 275328, 274944, 2117, 2325},
|
|
{"screenshot 3", 275049, 270336, 3926, 4183},
|
|
{"screenshot 4", 270431, 270208, 58, 949},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
info := claudeCacheReadTestInfo()
|
|
original := &dto.Usage{BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{
|
|
InputTokens: tc.prompt, CacheReadInputTokens: tc.cached, OutputTokens: tc.output,
|
|
})}
|
|
before, err := common.Marshal(original)
|
|
require.NoError(t, err)
|
|
canonical := effectiveBillingUsage(original)
|
|
adjusted, audit := normalizeClaudeCacheReadUsage(info, canonical)
|
|
require.NotNil(t, audit)
|
|
assert.Equal(t, "applied", audit.Status)
|
|
assert.Equal(t, tc.prompt, audit.OriginalInputTokens)
|
|
assert.Equal(t, tc.prompt-tc.cached, adjusted.PromptTokens)
|
|
assert.Equal(t, tc.prompt, adjusted.InputTokens)
|
|
assert.Equal(t, tc.prompt-tc.cached, adjusted.BillingUsage.ClaudeUsage.InputTokens)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
assert.Equal(t, tc.wantQuota, calculateTextQuotaSummary(ctx, info, adjusted).Quota)
|
|
params := BuildTieredTokenParams(adjusted, true, billingexpr.UsedVars("p * 2 + cr * 0.04 + c * 8"))
|
|
assert.Equal(t, float64(tc.prompt), params.Len)
|
|
assert.Equal(t, float64(tc.prompt-tc.cached), params.P)
|
|
cost, _, err := billingexpr.RunExprWithRequest("p * 2 + cr * 0.04 + c * 8", params, billingexpr.RequestInput{})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, tc.wantQuota, common.QuotaRound(cost/1_000_000*common.QuotaPerUnit*0.162))
|
|
withoutSeparateCache := BuildTieredTokenParams(adjusted, true, billingexpr.UsedVars("p * 2 + c * 8"))
|
|
assert.Equal(t, float64(tc.prompt), withoutSeparateCache.P)
|
|
again, _ := normalizeClaudeCacheReadUsage(info, effectiveBillingUsage(original))
|
|
assert.Equal(t, adjusted, again, "re-reading the same upstream snapshot must not accumulate corrections")
|
|
after, err := common.Marshal(original)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, before, after)
|
|
assert.Equal(t, tc.prompt, canonical.PromptTokens)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestClaudeCacheReadCompatibilityBoundaries(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
prompt, cached, creation int
|
|
estimated, disabled, openAI bool
|
|
model, returnedModel, status string
|
|
wantPrompt, wantInput int
|
|
}{
|
|
{name: "disabled standard Anthropic", prompt: 100, cached: 80, disabled: true, wantPrompt: 100, wantInput: 180},
|
|
{name: "another model", prompt: 100, cached: 80, model: "deepseek-v4-pro", wantPrompt: 100, wantInput: 180},
|
|
{name: "exact names only", prompt: 100, cached: 80, model: "deepseek-v4.1-flash-extra", wantPrompt: 100, wantInput: 180},
|
|
{name: "returned model cannot opt in", prompt: 100, cached: 80, model: "other", returnedModel: "deepseek-v4.1-flash", wantPrompt: 100, wantInput: 180},
|
|
{name: "returned model cannot opt out", prompt: 100, cached: 80, returnedModel: "other", status: "applied", wantPrompt: 20, wantInput: 100},
|
|
{name: "OpenAI snapshot remains authoritative", prompt: 100, cached: 80, openAI: true, wantPrompt: 100, wantInput: 100},
|
|
{name: "estimated usage", prompt: 100, cached: 80, estimated: true, wantPrompt: 100, wantInput: 180},
|
|
{name: "zero cache", prompt: 100, wantPrompt: 100, wantInput: 100},
|
|
{name: "full cache", prompt: 100, cached: 100, status: "applied", wantInput: 100},
|
|
{name: "cache writes remain separate", prompt: 100, cached: 80, creation: 10, status: "applied", wantPrompt: 20, wantInput: 110},
|
|
{name: "negative prompt", prompt: -1, cached: 80, status: "invalid_usage", wantPrompt: -1, wantInput: 79},
|
|
{name: "negative cache", prompt: 100, cached: -1, status: "invalid_usage", wantPrompt: 100, wantInput: 99},
|
|
{name: "cache exceeds input", prompt: 100, cached: 101, status: "invalid_usage", wantPrompt: 100, wantInput: 201},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
info := claudeCacheReadTestInfo()
|
|
if tc.disabled {
|
|
info.ChannelOtherSettings.ClaudeInputIncludesCacheReadModels = nil
|
|
}
|
|
if tc.model != "" {
|
|
info.UpstreamModelName = tc.model
|
|
}
|
|
if tc.returnedModel != "" {
|
|
info.ObserveResponseModel(tc.returnedModel)
|
|
info.UpstreamModelName = tc.returnedModel
|
|
}
|
|
usage := &dto.Usage{BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{
|
|
InputTokens: tc.prompt, CacheReadInputTokens: tc.cached, CacheCreationInputTokens: tc.creation, OutputTokens: 1,
|
|
})}
|
|
if tc.openAI {
|
|
usage.BillingUsage = dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: tc.prompt, CompletionTokens: 1,
|
|
PromptTokensDetails: dto.InputTokenDetails{CachedTokens: tc.cached}})
|
|
}
|
|
usage.BillingUsage.Estimated = tc.estimated
|
|
canonical := effectiveBillingUsage(usage)
|
|
adjusted, audit := normalizeClaudeCacheReadUsage(info, canonical)
|
|
assert.Equal(t, tc.wantPrompt, adjusted.PromptTokens)
|
|
assert.Equal(t, tc.wantInput, adjusted.InputTokens)
|
|
if tc.status == "" {
|
|
assert.Nil(t, audit)
|
|
assert.Same(t, canonical, adjusted)
|
|
} else {
|
|
require.NotNil(t, audit)
|
|
assert.Equal(t, tc.status, audit.Status)
|
|
}
|
|
assert.Equal(t, tc.creation, adjusted.PromptTokensDetails.CacheCreationTokensTotal())
|
|
})
|
|
}
|
|
}
|
|
|
|
func runClaudeCacheReadAccountingCases(t *testing.T, db, logDB *gorm.DB) {
|
|
t.Helper()
|
|
for _, stream := range []bool{false, true} {
|
|
for _, expression := range []bool{false, true} {
|
|
t.Run(fmt.Sprintf("stream=%t expression=%t", stream, expression), func(t *testing.T) {
|
|
const initialQuota = 1_000_000
|
|
info := claudeCacheReadTestInfo()
|
|
channel := model.Channel{Name: "cache-read-test", Key: "unused", Status: common.ChannelStatusEnabled}
|
|
channel.SetOtherSettings(info.ChannelOtherSettings)
|
|
require.NoError(t, db.Create(&channel).Error)
|
|
user := model.User{Username: fmt.Sprintf("cache_read_%t_%t", stream, expression), Quota: initialQuota, Status: common.UserStatusEnabled}
|
|
require.NoError(t, db.Create(&user).Error)
|
|
token := model.Token{UserId: user.Id, Key: fmt.Sprintf("cache-read-%t-%t", stream, expression), RemainQuota: initialQuota, Status: common.TokenStatusEnabled}
|
|
require.NoError(t, db.Create(&token).Error)
|
|
t.Cleanup(func() {
|
|
require.NoError(t, logDB.Where("user_id = ?", user.Id).Delete(&model.Log{}).Error)
|
|
require.NoError(t, db.Unscoped().Delete(&token).Error)
|
|
require.NoError(t, db.Unscoped().Delete(&user).Error)
|
|
require.NoError(t, db.Delete(&channel).Error)
|
|
})
|
|
require.NoError(t, db.First(&channel, channel.Id).Error)
|
|
settings := channel.GetOtherSettings()
|
|
settings.AllowSpeed = true
|
|
channel.SetOtherSettings(settings)
|
|
require.NoError(t, db.Model(&channel).Update("settings", channel.OtherSettings).Error)
|
|
require.NoError(t, db.First(&channel, channel.Id).Error)
|
|
assert.True(t, channel.GetOtherSettings().AllowSpeed)
|
|
assert.Equal(t, []string{"deepseek-v4.1-flash"}, channel.GetOtherSettings().ClaudeInputIncludesCacheReadModels)
|
|
info.ChannelOtherSettings = channel.GetOtherSettings()
|
|
info.ChannelId, info.UserId, info.TokenId, info.TokenKey = channel.Id, user.Id, token.Id, token.Key
|
|
info.UserGroup, info.UsingGroup = "default", "default"
|
|
info.UserSetting.BillingPreference = "wallet_only"
|
|
info.ForcePreConsume, info.IsStream = true, stream
|
|
if expression {
|
|
const expr = `tier("base", p * 2 + cr * 0.04 + c * 8)`
|
|
info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expr,
|
|
ExprHash: billingexpr.ExprHashString(expr), QuotaPerUnit: common.QuotaPerUnit, GroupRatio: 0.162}
|
|
}
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx.Request = httptest.NewRequest("POST", "/v1/messages", nil)
|
|
require.Nil(t, PreConsumeBilling(ctx, 50_000, info))
|
|
wireUsage := &dto.ClaudeUsage{InputTokens: 275328, CacheReadInputTokens: 274944, OutputTokens: 2117}
|
|
usage := &dto.Usage{BillingUsage: dto.NewClaudeMessagesBillingUsage(wireUsage)}
|
|
if stream {
|
|
state := &relayconvert.ClaudeResponseInfo{Usage: &dto.Usage{}}
|
|
relayconvert.FormatClaudeResponseInfo(&dto.ClaudeResponse{Type: "message_start", Message: &dto.ClaudeMediaMessage{Usage: &dto.ClaudeUsage{InputTokens: 275328}}}, nil, state)
|
|
relayconvert.FormatClaudeResponseInfo(&dto.ClaudeResponse{Type: "message_delta", Usage: wireUsage}, nil, state)
|
|
relayconvert.FinalizeClaudeStreamBillingUsage(state)
|
|
usage = state.Usage
|
|
}
|
|
PostTextConsumeQuota(ctx, info, usage, nil)
|
|
require.NoError(t, db.First(&user, user.Id).Error)
|
|
require.NoError(t, db.First(&token, token.Id).Error)
|
|
require.NoError(t, db.First(&channel, channel.Id).Error)
|
|
assert.Equal(t, initialQuota-2325, user.Quota)
|
|
assert.Equal(t, initialQuota-2325, token.RemainQuota)
|
|
assert.Equal(t, 2325, user.UsedQuota)
|
|
assert.Equal(t, int64(2325), channel.UsedQuota)
|
|
var log model.Log
|
|
require.NoError(t, logDB.Where("user_id = ?", user.Id).Take(&log).Error)
|
|
assert.Equal(t, 2325, log.Quota)
|
|
assert.Equal(t, 384, log.PromptTokens)
|
|
var other struct {
|
|
Admin struct {
|
|
Adjustment claudeCacheReadAdjustment `json:"cache_read_adjustment"`
|
|
} `json:"admin_info"`
|
|
}
|
|
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
|
assert.Equal(t, "applied", other.Admin.Adjustment.Status)
|
|
assert.Equal(t, 275328, other.Admin.Adjustment.OriginalInputTokens)
|
|
assert.Equal(t, 384, other.Admin.Adjustment.InputTokens)
|
|
assert.Equal(t, 275328, wireUsage.InputTokens)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) {
|
|
t.Helper()
|
|
const mixed = `len <= 32000 ? tier("short", fixed(0.01)) : tier("long", p * 2)`
|
|
const flat = `tier("request", fixed(0.01))`
|
|
const startingQuota = 2_000_000
|
|
const imageExpression = `tier("standard", p * 5 + cr * 1.25 + img * 8 + img_cr * 2 + c * 30)`
|
|
imageUsage := &dto.Usage{PromptTokens: 1000, CompletionTokens: 100, TotalTokens: 1100,
|
|
PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 300, ImageTokens: 600, CachedTokensDetails: &dto.CachedTokenDetails{ImageTokens: common.GetPointer(200)}}}
|
|
operation_setting.SetToolPriceForTest("fixed_billing_tool", 4)
|
|
t.Cleanup(func() { operation_setting.DeleteToolPriceForTest("fixed_billing_tool") })
|
|
for index, tc := range []struct {
|
|
name, expression string
|
|
estimate int
|
|
usage *dto.Usage
|
|
audio, stream, refund, insufficient, tool bool
|
|
realtime, reserveInsufficient bool
|
|
wallet, outboundImages int
|
|
groupRatio float64
|
|
want int
|
|
unit billingexpr.BillingUnit
|
|
requestedImages, actualImages int
|
|
}{
|
|
{name: "missing usage charges once", expression: flat, want: 5000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "zero usage charges once", expression: flat, usage: &dto.Usage{}, want: 5000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "stream charges once", expression: flat, stream: true, usage: &dto.Usage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120}, want: 5000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "audio zero usage charges once", expression: flat, audio: true, usage: &dto.Usage{}, want: 5000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "audio missing usage charges once", expression: flat, audio: true, want: 5000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "token reservation refunds to fixed price", expression: mixed, estimate: 50000, usage: &dto.Usage{PromptTokens: 100, TotalTokens: 100}, want: 5000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "fixed reservation settles token fallback", expression: mixed, estimate: 100, usage: &dto.Usage{PromptTokens: 50000, TotalTokens: 50000}, want: 50000, unit: billingexpr.BillingUnitToken},
|
|
{name: "missing usage uses estimated token fallback", expression: mixed, estimate: 50000, want: 50000, unit: billingexpr.BillingUnitToken},
|
|
{name: "evaluation error retains fixed reservation metadata", expression: `p == 50 ? tier("error", param("missing") * p + img_cr * 2) : tier("request", fixed(0.01))`, estimate: 100, usage: &dto.Usage{PromptTokens: 50, TotalTokens: 50}, want: 5000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "explicit zero remains free", expression: `tier("free", fixed(0))`, usage: &dto.Usage{PromptTokens: 100, TotalTokens: 100}, unit: billingexpr.BillingUnitRequest},
|
|
{name: "multipliers and separate tool surcharge", expression: flat + ` * (param("fast") == true ? 2 : 1)`, groupRatio: 1.5, tool: true, want: 18000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "failed request refunds exactly once", expression: flat, refund: true},
|
|
{name: "insufficient wallet never reserves tokens", expression: flat, insufficient: true},
|
|
{name: "image cache stream settles usage and refunds unused reservation", expression: imageExpression, estimate: 10000, usage: imageUsage, stream: true, want: 4113, unit: billingexpr.BillingUnitToken},
|
|
{name: "audio settlement records image cache billing inputs", expression: imageExpression, estimate: 10000, usage: imageUsage, audio: true, want: 4113, unit: billingexpr.BillingUnitToken},
|
|
{name: "realtime records actual expression inputs", expression: imageExpression, estimate: 10000, usage: imageUsage, realtime: true, want: 4000, unit: billingexpr.BillingUnitToken},
|
|
{name: "image cache insufficient wallet never reserves tokens", expression: imageExpression, estimate: 10000, insufficient: true},
|
|
{name: "image quantity refunds missing images", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 3, actualImages: 2, want: 40000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "image quantity keeps request when actual missing", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 3, want: 60000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "image quantity zero price stays free", expression: `tier("image", fixed(0)) * image_count`, requestedImages: 3, actualImages: 2, unit: billingexpr.BillingUnitRequest},
|
|
{name: "image quantity failure refunds reservation", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 3, refund: true},
|
|
{name: "image quantity exceeds one-image wallet before submission", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 4, wallet: 20000, insufficient: true},
|
|
{name: "image override reserves extra quantity", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 1, outboundImages: 4, want: 80000, unit: billingexpr.BillingUnitRequest},
|
|
{name: "image override cannot exceed remaining wallet", expression: `tier("image", fixed(0.04)) * image_count`, requestedImages: 1, outboundImages: 4, wallet: 40000, reserveInsufficient: true, refund: true},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
quota := startingQuota
|
|
if tc.insufficient {
|
|
quota = 1
|
|
}
|
|
if tc.wallet > 0 {
|
|
quota = tc.wallet
|
|
}
|
|
user := model.User{Username: fmt.Sprintf("fixed_billing_%d", index), Quota: quota, Status: common.UserStatusEnabled}
|
|
require.NoError(t, db.Create(&user).Error)
|
|
token := model.Token{UserId: user.Id, Key: fmt.Sprintf("fixed-billing-test-%d", index), Name: "fixed-billing", RemainQuota: startingQuota, Status: common.TokenStatusEnabled}
|
|
require.NoError(t, db.Create(&token).Error)
|
|
channel := model.Channel{Name: "fixed-billing", Key: "unused", Status: common.ChannelStatusEnabled}
|
|
require.NoError(t, db.Create(&channel).Error)
|
|
t.Cleanup(func() {
|
|
require.NoError(t, logDB.Where("user_id = ?", user.Id).Delete(&model.Log{}).Error)
|
|
require.NoError(t, db.Unscoped().Delete(&token).Error)
|
|
require.NoError(t, db.Unscoped().Delete(&user).Error)
|
|
require.NoError(t, db.Unscoped().Delete(&channel).Error)
|
|
})
|
|
group := tc.groupRatio
|
|
if group == 0 {
|
|
group = 1
|
|
}
|
|
request := &billingexpr.RequestInput{Body: []byte(`{"fast":true}`)}
|
|
if tc.requestedImages > 0 {
|
|
request.ImageCount = &tc.requestedImages
|
|
}
|
|
cost, trace, err := billingexpr.RunExprWithRequest(tc.expression, billingexpr.TokenParams{P: float64(tc.estimate), Len: float64(tc.estimate)}, *request)
|
|
require.NoError(t, err)
|
|
reservation, err := billingexpr.QuotaRoundStrict(cost / 1_000_000 * common.QuotaPerUnit * group)
|
|
require.NoError(t, err)
|
|
snapshot := &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: tc.expression, ExprHash: billingexpr.ExprHashString(tc.expression), QuotaPerUnit: common.QuotaPerUnit, GroupRatio: group, EstimatedTier: trace.MatchedTier, EstimatedBillingUnit: trace.BillingUnit, EstimatedFixedPrice: trace.FixedPrice, EstimatedQuotaAfterGroup: reservation}
|
|
snapshot.EstimatedImageCount = trace.ImageCount
|
|
info := &relaycommon.RelayInfo{UserId: user.Id, TokenId: token.Id, TokenKey: token.Key, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channel.Id}, OriginModelName: "fixed-test", UsingGroup: "default", UserGroup: "default", UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}, ForcePreConsume: true, StartTime: time.Now(), IsStream: tc.stream, RelayFormat: types.RelayFormatOpenAI, PriceData: hosttypes.PriceData{GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: group}}, TieredBillingSnapshot: snapshot, BillingRequestInput: request}
|
|
info.SetEstimatePromptTokens(tc.estimate)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
if tc.expression == imageExpression {
|
|
ctx.Request.URL.Path = "/v1/images/generations"
|
|
info.RelayMode = relayconstant.RelayModeImagesGenerations
|
|
}
|
|
apiErr := PreConsumeBilling(ctx, reservation, info)
|
|
if tc.insufficient {
|
|
require.NotNil(t, apiErr)
|
|
assert.Equal(t, types.ErrorCodeInsufficientUserQuota, apiErr.GetErrorCode())
|
|
} else {
|
|
require.Nil(t, apiErr)
|
|
held, err := model.GetUserQuota(user.Id, true)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, quota-reservation, held)
|
|
if tc.outboundImages > 0 {
|
|
reserveErr := PrepareImageBillingForRequest(ctx, info, tc.outboundImages)
|
|
if tc.reserveInsufficient {
|
|
require.NotNil(t, reserveErr)
|
|
assert.Equal(t, types.ErrorCodeInsufficientUserQuota, reserveErr.GetErrorCode())
|
|
assert.Equal(t, reservation, info.Billing.GetPreConsumedQuota())
|
|
} else {
|
|
require.Nil(t, reserveErr)
|
|
assert.Equal(t, tc.want, info.Billing.GetPreConsumedQuota())
|
|
}
|
|
}
|
|
if tc.refund {
|
|
refunded := make(chan struct{}, 1)
|
|
const callback = "fixed_billing_refund_observed"
|
|
require.NoError(t, db.Callback().Update().After("gorm:commit_or_rollback_transaction").Register(callback, func(tx *gorm.DB) {
|
|
if tx.Statement.Table == "tokens" && tx.Error == nil {
|
|
select {
|
|
case refunded <- struct{}{}:
|
|
default:
|
|
}
|
|
}
|
|
}))
|
|
t.Cleanup(func() { require.NoError(t, db.Callback().Update().Remove(callback)) })
|
|
info.Billing.Refund(ctx)
|
|
info.Billing.Refund(ctx)
|
|
select {
|
|
case <-refunded:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("refund did not finish")
|
|
}
|
|
} else {
|
|
info.UpdateImageCount(int64(tc.actualImages))
|
|
if tc.tool {
|
|
info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{"fixed_billing_tool": {CallCount: 1}}}
|
|
}
|
|
if tc.realtime {
|
|
PostWssConsumeQuota(ctx, info, info.OriginModelName, &dto.RealtimeUsage{
|
|
InputTokens: tc.usage.PromptTokens, OutputTokens: tc.usage.CompletionTokens, TotalTokens: tc.usage.TotalTokens,
|
|
}, nil)
|
|
} else if tc.audio {
|
|
PostAudioConsumeQuota(ctx, info, tc.usage, nil)
|
|
} else {
|
|
PostTextConsumeQuota(ctx, info, tc.usage, nil)
|
|
}
|
|
require.NoError(t, info.Billing.Settle(tc.want), "a repeated settlement must not charge again")
|
|
var log model.Log
|
|
require.NoError(t, logDB.Where("user_id = ?", user.Id).Take(&log).Error)
|
|
assert.Equal(t, tc.want, log.Quota)
|
|
assert.Equal(t, tc.stream, log.IsStream)
|
|
assert.NotContains(t, log.Content, "nothing was charged")
|
|
var other map[string]any
|
|
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
|
assert.Equal(t, string(tc.unit), other["billing_unit"])
|
|
if tc.requestedImages > 0 {
|
|
count := tc.actualImages
|
|
if count == 0 {
|
|
count = tc.requestedImages
|
|
if tc.outboundImages > 0 {
|
|
count = tc.outboundImages
|
|
}
|
|
}
|
|
assert.Equal(t, float64(count), other["image_count"])
|
|
assert.Equal(t, tc.requestedImages, *request.ImageCount, "actual count must not mutate the frozen request")
|
|
}
|
|
if tc.expression == imageExpression {
|
|
billable, ok := other["billing_tokens"].(map[string]any)
|
|
require.True(t, ok)
|
|
if tc.realtime {
|
|
assert.Equal(t, float64(0), other["image_cache_tokens"])
|
|
assert.Equal(t, float64(1000), billable["p"])
|
|
assert.Equal(t, float64(0), billable["cr"])
|
|
assert.Equal(t, float64(0), billable["img"])
|
|
} else {
|
|
if !tc.audio {
|
|
assert.Equal(t, float64(300), other["cache_tokens"])
|
|
}
|
|
assert.Equal(t, float64(200), other["image_cache_tokens"])
|
|
assert.Equal(t, float64(300), billable["p"])
|
|
assert.Equal(t, float64(100), billable["cr"])
|
|
assert.Equal(t, float64(400), billable["img"])
|
|
}
|
|
} else {
|
|
assert.NotContains(t, other, "billing_tokens")
|
|
assert.NotContains(t, other, "image_cache_tokens")
|
|
}
|
|
if tc.unit == billingexpr.BillingUnitRequest {
|
|
assert.Contains(t, other, "fixed_price")
|
|
} else {
|
|
assert.NotContains(t, other, "fixed_price")
|
|
}
|
|
}
|
|
}
|
|
require.NoError(t, db.First(&user, user.Id).Error)
|
|
require.NoError(t, db.First(&token, token.Id).Error)
|
|
assert.Equal(t, quota-tc.want, user.Quota)
|
|
assert.Equal(t, startingQuota-tc.want, token.RemainQuota)
|
|
assert.Equal(t, tc.want, user.UsedQuota)
|
|
assert.Equal(t, tc.want, token.UsedQuota)
|
|
if !tc.refund && !tc.insufficient {
|
|
assert.Equal(t, 1, user.RequestCount)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryUnifiedForClaudeSemantic(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 1000,
|
|
CompletionTokens: 200,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 100,
|
|
CachedCreationTokens: 50,
|
|
},
|
|
ClaudeCacheCreation5mTokens: 10,
|
|
ClaudeCacheCreation1hTokens: 20,
|
|
}
|
|
|
|
priceData := hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25,
|
|
CacheCreation5mRatio: 1.25,
|
|
CacheCreation1hRatio: 2,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{
|
|
GroupRatio: 1,
|
|
},
|
|
}
|
|
|
|
chatRelayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
FinalRequestRelayFormat: types.RelayFormatClaude,
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: priceData,
|
|
StartTime: time.Now(),
|
|
}
|
|
messageRelayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatClaude,
|
|
FinalRequestRelayFormat: types.RelayFormatClaude,
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: priceData,
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
chatSummary := calculateTextQuotaSummary(ctx, chatRelayInfo, usage)
|
|
messageSummary := calculateTextQuotaSummary(ctx, messageRelayInfo, usage)
|
|
|
|
require.Equal(t, messageSummary.Quota, chatSummary.Quota)
|
|
require.Equal(t, messageSummary.CacheCreationTokens5m, chatSummary.CacheCreationTokens5m)
|
|
require.Equal(t, messageSummary.CacheCreationTokens1h, chatSummary.CacheCreationTokens1h)
|
|
require.True(t, chatSummary.IsClaudeUsageSemantic)
|
|
require.Equal(t, 1488, chatSummary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryUsesSplitClaudeCacheCreationRatios(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
FinalRequestRelayFormat: types.RelayFormatClaude,
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
CacheRatio: 0,
|
|
CacheCreationRatio: 1,
|
|
CacheCreation5mRatio: 2,
|
|
CacheCreation1hRatio: 3,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{
|
|
GroupRatio: 1,
|
|
},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 0,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedCreationTokens: 10,
|
|
},
|
|
ClaudeCacheCreation5mTokens: 2,
|
|
ClaudeCacheCreation1hTokens: 3,
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
// 100 + remaining(5)*1 + 2*2 + 3*3 = 118
|
|
require.Equal(t, 118, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryUsesAnthropicUsageSemanticFromUpstreamUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25,
|
|
CacheCreation5mRatio: 1.25,
|
|
CacheCreation1hRatio: 2,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{
|
|
GroupRatio: 1,
|
|
},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 1000,
|
|
CompletionTokens: 200,
|
|
UsageSemantic: "anthropic",
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 100,
|
|
CachedCreationTokens: 50,
|
|
},
|
|
ClaudeCacheCreation5mTokens: 10,
|
|
ClaudeCacheCreation1hTokens: 20,
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
require.True(t, summary.IsClaudeUsageSemantic)
|
|
require.Equal(t, "anthropic", summary.UsageSemantic)
|
|
require.Equal(t, 1488, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryUsesClaudeBillingUsageBeforeTopLevelUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25,
|
|
CacheCreation5mRatio: 1.25,
|
|
CacheCreation1hRatio: 2,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 999,
|
|
CompletionTokens: 999,
|
|
TotalTokens: 1998,
|
|
BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{
|
|
InputTokens: 70,
|
|
CacheReadInputTokens: 30,
|
|
CacheCreationInputTokens: 20,
|
|
OutputTokens: 7,
|
|
CacheCreation: &dto.ClaudeCacheCreationUsage{
|
|
Ephemeral5mInputTokens: 12,
|
|
Ephemeral1hInputTokens: 8,
|
|
},
|
|
}),
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveBillingUsage(usage))
|
|
|
|
require.True(t, summary.IsClaudeUsageSemantic)
|
|
require.Equal(t, dto.BillingUsageSemanticAnthropic, summary.UsageSemantic)
|
|
require.Equal(t, 70, summary.PromptTokens)
|
|
require.Equal(t, 7, summary.CompletionTokens)
|
|
require.Equal(t, 30, summary.CacheTokens)
|
|
require.Equal(t, 20, summary.CacheCreationTokens)
|
|
require.Equal(t, 12, summary.CacheCreationTokens5m)
|
|
require.Equal(t, 8, summary.CacheCreationTokens1h)
|
|
require.Equal(t, 118, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryUsesGeminiBillingUsageBeforeTopLevelUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "gemini-2.5-flash",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
CacheRatio: 0.1,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 999,
|
|
CompletionTokens: 999,
|
|
TotalTokens: 1998,
|
|
BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{
|
|
PromptTokenCount: 100,
|
|
ToolUsePromptTokenCount: 5,
|
|
CandidatesTokenCount: 20,
|
|
ThoughtsTokenCount: 3,
|
|
TotalTokenCount: 128,
|
|
CachedContentTokenCount: 7,
|
|
}),
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveBillingUsage(usage))
|
|
|
|
require.False(t, summary.IsClaudeUsageSemantic)
|
|
require.Equal(t, dto.BillingUsageSemanticGemini, summary.UsageSemantic)
|
|
require.Equal(t, 105, summary.PromptTokens)
|
|
require.Equal(t, 23, summary.CompletionTokens)
|
|
require.Equal(t, 7, summary.CacheTokens)
|
|
require.Equal(t, 128, summary.TotalTokens)
|
|
require.Equal(t, 145, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatClaude,
|
|
OriginModelName: "gpt-4o",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 999,
|
|
CompletionTokens: 999,
|
|
TotalTokens: 1998,
|
|
BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{
|
|
PromptTokens: 80,
|
|
CompletionTokens: 9,
|
|
TotalTokens: 89,
|
|
}),
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveBillingUsage(usage))
|
|
|
|
require.False(t, summary.IsClaudeUsageSemantic)
|
|
require.Equal(t, dto.BillingUsageSemanticOpenAI, summary.UsageSemantic)
|
|
require.Equal(t, 80, summary.PromptTokens)
|
|
require.Equal(t, 9, summary.CompletionTokens)
|
|
require.Equal(t, 89, summary.TotalTokens)
|
|
require.Equal(t, 98, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryUsesOpenAIResponsesInputTokenDetails(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "gpt-4o",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
CacheRatio: 0.25,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
responsesUsage := &dto.Usage{
|
|
InputTokens: 100,
|
|
OutputTokens: 10,
|
|
TotalTokens: 110,
|
|
InputTokensDetails: &dto.InputTokenDetails{
|
|
CachedTokens: 40,
|
|
},
|
|
}
|
|
convertedUsage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 10,
|
|
TotalTokens: 110,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 40,
|
|
},
|
|
BillingUsage: dto.NewOpenAIResponsesBillingUsage(responsesUsage),
|
|
}
|
|
|
|
effectiveUsage := effectiveBillingUsage(convertedUsage)
|
|
require.Equal(t, 40, effectiveUsage.PromptTokensDetails.CachedTokens)
|
|
require.Zero(t, convertedUsage.BillingUsage.OpenAIUsage.PromptTokensDetails.CachedTokens)
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveUsage)
|
|
require.Equal(t, 40, summary.CacheTokens)
|
|
// 60 uncached input + 40*0.25 cached input + 10*2 output = 90.
|
|
require.Equal(t, 90, summary.Quota)
|
|
}
|
|
|
|
func TestUsageFromOpenAIBillingUsageNormalizesCacheDetailsWithoutOverwritingCanonicalValues(t *testing.T) {
|
|
responsesUsage := &dto.Usage{
|
|
InputTokens: 100,
|
|
OutputTokens: 10,
|
|
PromptCacheHitTokens: 55,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 8,
|
|
TextTokens: 12,
|
|
},
|
|
InputTokensDetails: &dto.InputTokenDetails{
|
|
CachedTokens: 40,
|
|
CachedCreationTokens: 5,
|
|
CacheWriteTokens: 6,
|
|
TextTokens: 60,
|
|
ImageTokens: 7,
|
|
AudioTokens: 9,
|
|
},
|
|
}
|
|
|
|
billingUsage := dto.NewOpenAIResponsesBillingUsage(responsesUsage)
|
|
usage := effectiveBillingUsage(&dto.Usage{BillingUsage: billingUsage})
|
|
|
|
require.Equal(t, 8, usage.PromptTokensDetails.CachedTokens)
|
|
require.Equal(t, 5, usage.PromptTokensDetails.CachedCreationTokens)
|
|
require.Equal(t, 6, usage.PromptTokensDetails.CacheWriteTokens)
|
|
require.Equal(t, 12, usage.PromptTokensDetails.TextTokens)
|
|
require.Equal(t, 7, usage.PromptTokensDetails.ImageTokens)
|
|
require.Equal(t, 9, usage.PromptTokensDetails.AudioTokens)
|
|
require.Zero(t, billingUsage.OpenAIUsage.PromptTokensDetails.CachedCreationTokens)
|
|
}
|
|
|
|
func TestUsageFromOpenAIBillingUsageFallsBackToPromptCacheHitTokens(t *testing.T) {
|
|
usage := effectiveBillingUsage(&dto.Usage{
|
|
BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 10,
|
|
PromptCacheHitTokens: 35,
|
|
}),
|
|
})
|
|
|
|
require.Equal(t, 35, usage.PromptTokensDetails.CachedTokens)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryNormalizesOpenAIResponsesBillingUsageDetails(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatClaude,
|
|
OriginModelName: "gpt-5.6-sol",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
CacheRatio: 0.5,
|
|
CacheCreationRatio: 2,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
responsesDetails := dto.InputTokenDetails{
|
|
CachedTokens: 80,
|
|
CacheWriteTokens: 10,
|
|
TextTokens: 100,
|
|
}
|
|
usage := &dto.Usage{
|
|
PromptTokens: 999,
|
|
CompletionTokens: 999,
|
|
BillingUsage: dto.NewOpenAIResponsesBillingUsage(&dto.Usage{
|
|
InputTokens: 100,
|
|
OutputTokens: 10,
|
|
TotalTokens: 110,
|
|
InputTokensDetails: &responsesDetails,
|
|
}),
|
|
}
|
|
|
|
effectiveUsage := effectiveBillingUsage(usage)
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveUsage)
|
|
|
|
require.Equal(t, dto.BillingUsageSourceOAIResponses, effectiveUsage.UsageSource)
|
|
require.Equal(t, responsesDetails, effectiveUsage.PromptTokensDetails)
|
|
require.Equal(t, 100, summary.PromptTokens)
|
|
require.Equal(t, 10, summary.CompletionTokens)
|
|
require.Equal(t, 80, summary.CacheTokens)
|
|
require.Equal(t, 10, summary.CacheCreationTokens)
|
|
// (100-80-10) + 80*0.5 + 10*2 + 10*2 = 90
|
|
require.Equal(t, 90, summary.Quota)
|
|
}
|
|
|
|
func TestUsageBillingPathForLog(t *testing.T) {
|
|
require.Equal(t, usageBillingPathAnthropic, usageBillingPathForLog(true, &dto.Usage{
|
|
BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}),
|
|
}))
|
|
invalidBillingUsage := &dto.Usage{
|
|
PromptTokens: 1,
|
|
BillingUsage: &dto.BillingUsage{
|
|
Source: dto.BillingUsageSourceClaudeMessages,
|
|
Semantic: dto.BillingUsageSemanticAnthropic,
|
|
},
|
|
}
|
|
require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, invalidBillingUsage))
|
|
require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, invalidBillingUsage))
|
|
require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, &dto.Usage{}))
|
|
require.Equal(t, usageBillingPathOpenAI, usageBillingPathForLog(false, &dto.Usage{
|
|
BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: 1}),
|
|
}))
|
|
require.Equal(t, usageBillingPathAnthropic, usageBillingPathForLog(false, &dto.Usage{
|
|
BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}),
|
|
}))
|
|
require.Equal(t, usageBillingPathGemini, usageBillingPathForLog(false, &dto.Usage{
|
|
BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{PromptTokenCount: 1}),
|
|
}))
|
|
require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(true, &dto.Usage{
|
|
BillingUsage: dto.NewEstimatedGeminiChatBillingUsage(&dto.Usage{PromptTokens: 1}),
|
|
}))
|
|
}
|
|
|
|
func TestAppendUsageBillingPathForLogWritesAdminInfo(t *testing.T) {
|
|
other := model.NewLogOther()
|
|
appendUsageBillingPathForLog(other, true, &dto.Usage{
|
|
BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}),
|
|
})
|
|
|
|
adminInfo, ok := other.Snapshot()["admin_info"].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, usageBillingPathAnthropic, adminInfo["usage_billing_path"])
|
|
|
|
other = model.NewLogOther()
|
|
appendUsageBillingPathForLog(other, true, nil)
|
|
adminInfo, ok = other.Snapshot()["admin_info"].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, usageBillingPathLocal, adminInfo["usage_billing_path"])
|
|
}
|
|
|
|
func TestCacheWriteTokensTotal(t *testing.T) {
|
|
t.Run("split cache creation", func(t *testing.T) {
|
|
summary := textQuotaSummary{
|
|
CacheCreationTokens: 50,
|
|
CacheCreationTokens5m: 10,
|
|
CacheCreationTokens1h: 20,
|
|
}
|
|
require.Equal(t, 50, cacheWriteTokensTotal(summary))
|
|
})
|
|
|
|
t.Run("legacy cache creation", func(t *testing.T) {
|
|
summary := textQuotaSummary{CacheCreationTokens: 50}
|
|
require.Equal(t, 50, cacheWriteTokensTotal(summary))
|
|
})
|
|
|
|
t.Run("split cache creation without aggregate remainder", func(t *testing.T) {
|
|
summary := textQuotaSummary{
|
|
CacheCreationTokens5m: 10,
|
|
CacheCreationTokens1h: 20,
|
|
}
|
|
require.Equal(t, 30, cacheWriteTokensTotal(summary))
|
|
})
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryHandlesLegacyClaudeDerivedOpenAIUsage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 5,
|
|
CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25,
|
|
CacheCreation5mRatio: 1.25,
|
|
CacheCreation1hRatio: 2,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 62,
|
|
CompletionTokens: 95,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 3544,
|
|
},
|
|
ClaudeCacheCreation5mTokens: 586,
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
// 62 + 3544*0.1 + 586*1.25 + 95*5 = 1624.9 => 1624
|
|
require.Equal(t, 1624, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryBillsOpenAICacheWriteTokens(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatOpenAI,
|
|
OriginModelName: "gpt-5.1",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 2,
|
|
CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
t.Run("uncached remainder stays positive", func(t *testing.T) {
|
|
usage := &dto.Usage{
|
|
PromptTokens: 1473,
|
|
CompletionTokens: 19,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CacheWriteTokens: 1470,
|
|
},
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
require.Equal(t, 1470, summary.CacheCreationTokens)
|
|
// (1473-0-1470) + 1470*1.25 + 19*2 = 3 + 1837.5 + 38 = 1878.5 => 1879
|
|
require.Equal(t, 1879, summary.Quota)
|
|
})
|
|
|
|
t.Run("uncached remainder clamps to zero", func(t *testing.T) {
|
|
// Real OpenAI payload shape: cached_tokens + cache_write_tokens exceeds
|
|
// prompt_tokens because both are unadjusted prefix counts. The negative
|
|
// remainder must clamp to zero, never turn into a negative base charge.
|
|
usage := &dto.Usage{
|
|
PromptTokens: 3619,
|
|
CompletionTokens: 36,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 2921,
|
|
CacheWriteTokens: 3616,
|
|
},
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
require.Equal(t, 3619, summary.PromptTokens)
|
|
require.Equal(t, 3616, summary.CacheCreationTokens)
|
|
// max(3619-2921-3616, 0) + 2921*0.1 + 3616*1.25 + 36*2 = 4884.1 => 4884
|
|
require.Equal(t, 4884, summary.Quota)
|
|
})
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummarySeparatesOpenRouterCacheReadFromPromptBilling(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "openai/gpt-4.1",
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelType: constant.ChannelTypeOpenRouter,
|
|
},
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 2604,
|
|
CompletionTokens: 383,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 2432,
|
|
},
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
// OpenRouter OpenAI-format display keeps prompt_tokens as total input,
|
|
// but billing still separates normal input from cache read tokens.
|
|
// quota = (2604 - 2432) + 2432*0.1 + 383 = 798.2 => 798
|
|
require.Equal(t, 2604, summary.PromptTokens)
|
|
require.Equal(t, 798, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummarySeparatesOpenRouterCacheCreationFromPromptBilling(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "openai/gpt-4.1",
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelType: constant.ChannelTypeOpenRouter,
|
|
},
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
CacheCreationRatio: 1.25,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 2604,
|
|
CompletionTokens: 383,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedCreationTokens: 100,
|
|
},
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
// prompt_tokens is still logged as total input, but cache creation is billed separately.
|
|
// quota = (2604 - 100) + 100*1.25 + 383 = 3012
|
|
require.Equal(t, 2604, summary.PromptTokens)
|
|
require.Equal(t, 3012, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryKeepsPrePRClaudeOpenRouterBilling(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
FinalRequestRelayFormat: types.RelayFormatClaude,
|
|
OriginModelName: "anthropic/claude-3.7-sonnet",
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ChannelType: constant.ChannelTypeOpenRouter,
|
|
},
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 2604,
|
|
CompletionTokens: 383,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 2432,
|
|
},
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
// Pre-PR PostClaudeConsumeQuota behavior for OpenRouter:
|
|
// prompt = 2604 - 2432 = 172
|
|
// quota = 172 + 2432*0.1 + 383 = 798.2 => 798
|
|
require.True(t, summary.IsClaudeUsageSemantic)
|
|
require.Equal(t, 172, summary.PromptTokens)
|
|
require.Equal(t, 798, summary.Quota)
|
|
}
|
|
|
|
func TestComposeTieredTextQuotaKeepsToolCallSurcharges(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
// 11 $/1K => 0.011 per completed image output, matching the prior fixed low-tier charge.
|
|
operation_setting.SetToolPriceForTest(dto.BuildInToolImageGeneration, 11.0)
|
|
t.Cleanup(func() {
|
|
operation_setting.DeleteToolPriceForTest(dto.BuildInToolImageGeneration)
|
|
})
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "o1",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolWebSearchPreview: {
|
|
CallCount: 1,
|
|
},
|
|
dto.BuildInToolFileSearch: {
|
|
CallCount: 2,
|
|
},
|
|
dto.BuildInToolImageGeneration: {
|
|
CallCount: 1,
|
|
},
|
|
},
|
|
},
|
|
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
|
BillingMode: "tiered_expr",
|
|
GroupRatio: 1,
|
|
EstimatedQuotaBeforeGroup: 1000,
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 50,
|
|
TotalTokens: 150,
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
quota := composeTieredTextQuota(relayInfo, summary, 1000, &billingexpr.TieredResult{
|
|
ActualQuotaBeforeGroup: 1000,
|
|
ActualQuotaAfterGroup: 1000,
|
|
})
|
|
|
|
require.Equal(t, int64(13000), summary.ToolCallSurchargeQuota.Round(0).IntPart())
|
|
require.Equal(t, 14000, quota)
|
|
}
|
|
|
|
func TestComposeTieredTextQuotaFallbackKeepsToolCallSurcharges(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1.25},
|
|
},
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolWebSearch: {ToolName: dto.BuildInToolWebSearch, CallCount: 2},
|
|
},
|
|
},
|
|
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
|
BillingMode: "tiered_expr",
|
|
GroupRatio: 1.25,
|
|
EstimatedQuotaBeforeGroup: 1000,
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 50,
|
|
TotalTokens: 150,
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
quota := composeTieredTextQuota(relayInfo, summary, 1250, nil)
|
|
|
|
require.Equal(t, int64(12500), summary.ToolCallSurchargeQuota.Round(0).IntPart())
|
|
require.Equal(t, 13750, quota)
|
|
}
|
|
|
|
func TestComposeTieredTextQuotaErrorFallbackUsesPreConsumedQuota(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(w)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1.25},
|
|
},
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolWebSearch: {ToolName: dto.BuildInToolWebSearch, CallCount: 2},
|
|
},
|
|
},
|
|
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
|
BillingMode: "tiered_expr",
|
|
GroupRatio: 1.25,
|
|
EstimatedQuotaBeforeGroup: 1000,
|
|
},
|
|
StartTime: time.Now(),
|
|
}
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 50,
|
|
TotalTokens: 150,
|
|
}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
// tieredResult=nil simulates a settlement error where TryTieredSettle
|
|
// falls back to FinalPreConsumedQuota (2000), which differs from
|
|
// EstimatedQuotaBeforeGroup * GroupRatio (1250).
|
|
preConsumedFallback := 2000
|
|
quota := composeTieredTextQuota(relayInfo, summary, preConsumedFallback, nil)
|
|
|
|
require.Equal(t, int64(12500), summary.ToolCallSurchargeQuota.Round(0).IntPart())
|
|
require.Equal(t, 14500, quota)
|
|
}
|
|
|
|
// TestTryTieredSettleRecordsClampOnOverflow guards that an oversized tiered
|
|
// settlement both saturates the quota and records the clamp on RelayInfo, so
|
|
// every consume path (text, audio, WSS) can surface it under admin_info.
|
|
func TestTryTieredSettleRecordsClampOnOverflow(t *testing.T) {
|
|
// exprOutput = p * 1e12; quotaBeforeGroup = p*1e12 / 1e6 * 5e5 far exceeds
|
|
// the supported single-request range and must saturate.
|
|
exprStr := `tier("base", p * 1000000000000)`
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "overflow-model",
|
|
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
|
BillingMode: "tiered_expr",
|
|
ExprString: exprStr,
|
|
ExprHash: billingexpr.ExprHashString(exprStr),
|
|
GroupRatio: 1,
|
|
QuotaPerUnit: 500_000,
|
|
},
|
|
}
|
|
|
|
ok, quota, result := TryTieredSettle(relayInfo, billingexpr.TokenParams{P: 1_000_000_000})
|
|
|
|
require.True(t, ok)
|
|
require.NotNil(t, result)
|
|
require.Equal(t, common.MaxQuota, quota, "oversized settlement must clamp, never wrap negative")
|
|
require.NotNil(t, relayInfo.QuotaClamp, "clamp must be recorded on RelayInfo for admin auditing")
|
|
require.Equal(t, common.QuotaClampOverflow, relayInfo.QuotaClamp.Kind)
|
|
}
|
|
|
|
// TestTryTieredSettleNoClampInRange confirms an in-range settlement leaves
|
|
// RelayInfo.QuotaClamp nil.
|
|
func TestTryTieredSettleNoClampInRange(t *testing.T) {
|
|
exprStr := `tier("base", p * 2 + c * 10)`
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "in-range-model",
|
|
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
|
BillingMode: "tiered_expr",
|
|
ExprString: exprStr,
|
|
ExprHash: billingexpr.ExprHashString(exprStr),
|
|
GroupRatio: 1,
|
|
QuotaPerUnit: 500_000,
|
|
},
|
|
}
|
|
|
|
ok, _, result := TryTieredSettle(relayInfo, billingexpr.TokenParams{P: 1000, C: 500})
|
|
|
|
require.True(t, ok)
|
|
require.NotNil(t, result)
|
|
require.Nil(t, relayInfo.QuotaClamp, "in-range settlement must not record a clamp")
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryFixedPriceAppliesImageCountOnceAndAllowsOverride(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
priceData := hosttypes.PriceData{
|
|
ModelPrice: 0.12,
|
|
UsePrice: true,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{
|
|
GroupRatio: 1,
|
|
},
|
|
}
|
|
priceData.AddOtherRatio("n", 3)
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "dall-e-3",
|
|
PriceData: priceData,
|
|
StartTime: time.Now(),
|
|
}
|
|
usage := &dto.Usage{PromptTokens: 1, TotalTokens: 1}
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
require.Equal(t, 180000, summary.Quota)
|
|
|
|
// An adaptor-reported actual count replaces the requested count rather
|
|
// than multiplying it a second time.
|
|
relayInfo.PriceData.AddOtherRatio("n", 2)
|
|
summary = calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
require.Equal(t, 120000, summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextToolCallSurchargeGeneralizedBuiltInTools(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
|
|
operation_setting.SetToolPriceForTest("my_fn", 5.0)
|
|
t.Cleanup(func() {
|
|
operation_setting.DeleteToolPriceForTest("my_fn")
|
|
})
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "o1",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolWebSearchPreview: {CallCount: 2},
|
|
"my_fn": {CallCount: 3},
|
|
"unpriced": {CallCount: 5},
|
|
},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{
|
|
ModelName: "o1",
|
|
GroupRatio: 1,
|
|
}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
expected := decimal.NewFromFloat((10.0*2 + 5.0*3) / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expected.Equal(surcharge), "got %s want %s", surcharge, expected)
|
|
require.Len(t, summary.ToolSurchargeItems, 2)
|
|
assert.Equal(t, "my_fn", summary.ToolSurchargeItems[0].Name)
|
|
assert.Equal(t, 3, summary.ToolSurchargeItems[0].Count)
|
|
assert.Equal(t, 5.0, summary.ToolSurchargeItems[0].Price)
|
|
assert.Equal(t, dto.BuildInToolWebSearchPreview, summary.ToolSurchargeItems[1].Name)
|
|
assert.Equal(t, 2, summary.ToolSurchargeItems[1].Count)
|
|
assert.Equal(t, 10.0, summary.ToolSurchargeItems[1].Price)
|
|
}
|
|
|
|
func TestCalculateTextToolCallSurchargeKeepsSearchPreviewFallbackWithCustomFunctions(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
|
|
operation_setting.SetToolPriceForTest("my_fn", 5)
|
|
t.Cleanup(func() {
|
|
operation_setting.DeleteToolPriceForTest("my_fn")
|
|
})
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayMode: relayconstant.RelayModeChatCompletions,
|
|
OriginModelName: "gpt-4o-search-preview",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
"my_fn": {CallCount: 1},
|
|
},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{
|
|
ModelName: relayInfo.OriginModelName,
|
|
GroupRatio: 1,
|
|
}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
|
|
require.Len(t, summary.ToolSurchargeItems, 2)
|
|
assert.Equal(t, "my_fn", summary.ToolSurchargeItems[0].Name)
|
|
assert.Equal(t, dto.BuildInToolWebSearchPreview, summary.ToolSurchargeItems[1].Name)
|
|
expected := decimal.NewFromFloat((5.0 + 25.0) / 1000).
|
|
Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expected.Equal(surcharge), "got %s want %s", surcharge, expected)
|
|
}
|
|
|
|
func TestCalculateTextToolCallSurchargeDoesNotInferSearchForResponses(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
RelayMode: relayconstant.RelayModeResponses,
|
|
OriginModelName: "gpt-4o-search-preview",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{
|
|
ModelName: relayInfo.OriginModelName,
|
|
GroupRatio: 1,
|
|
}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
|
|
assert.True(t, surcharge.IsZero())
|
|
assert.Empty(t, summary.ToolSurchargeItems)
|
|
}
|
|
|
|
// Every BuiltInTools count is bounded by MaxBillableToolCallCount at
|
|
// settlement, whatever recorded it (upstream-reported counts and per-item
|
|
// counts alike).
|
|
func TestCalculateTextToolCallSurchargeClampsCallCount(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
|
|
for recorded, billed := range map[int]int{5: 5, 10_000: 10_000, 10_005: 10_000} {
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "claude-3-7-sonnet",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolWebSearch: {CallCount: recorded},
|
|
},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{ModelName: relayInfo.OriginModelName, GroupRatio: 1}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
|
|
assert.Equal(t, []ToolSurchargeItem{{Name: dto.BuildInToolWebSearch, Count: billed, Price: 10}}, summary.ToolSurchargeItems, recorded)
|
|
expected := decimal.NewFromFloat(10.0 * float64(billed) / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expected.Equal(surcharge), "recorded %d: got %s want %s", recorded, surcharge, expected)
|
|
}
|
|
}
|
|
|
|
func TestMergeToolSurchargeItemsSaturatesCountOverflow(t *testing.T) {
|
|
items := []ToolSurchargeItem{
|
|
{Name: "custom_fn", Count: math.MaxInt, Price: 5},
|
|
{Name: "custom_fn", Count: 1, Price: 5},
|
|
}
|
|
|
|
merged := mergeToolSurchargeItems(items)
|
|
|
|
require.Len(t, merged, 1)
|
|
assert.Equal(t, math.MaxInt, merged[0].Count)
|
|
}
|
|
|
|
// A zero-token request (e.g. /v1/alpha/search returns no usage) must still
|
|
// bill a tool-call surcharge. Regression for the TotalTokens==0 gate zeroing
|
|
// out the surcharge quota.
|
|
func TestCalculateTextQuotaSummaryZeroTokensStillBillsToolSurcharge(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "o1",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolWebSearchPreview: {CallCount: 1},
|
|
},
|
|
},
|
|
}
|
|
relayInfo.PriceData.GroupRatioInfo.GroupRatio = 1
|
|
|
|
usage := &dto.Usage{} // zero tokens, mirrors alpha search
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
require.Equal(t, 0, summary.TotalTokens)
|
|
assert.False(t, summary.ToolCallSurchargeQuota.IsZero(), "surcharge should be computed")
|
|
assert.Greater(t, summary.Quota, 0, "quota must not be zeroed out for a zero-token web search request")
|
|
expected := common.QuotaFromDecimal(summary.ToolCallSurchargeQuota)
|
|
assert.Equal(t, expected, summary.Quota)
|
|
}
|
|
|
|
func TestClaudeCacheOnlyUsageRemainsBillable(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
for _, tc := range []struct {
|
|
name string
|
|
usage dto.ClaudeUsage
|
|
wantQuota int
|
|
}{
|
|
{name: "one-hour cache write", usage: dto.ClaudeUsage{CacheCreationInputTokens: 20000, CacheCreation: &dto.ClaudeCacheCreationUsage{Ephemeral1hInputTokens: 20000}}, wantQuota: 100000},
|
|
{name: "five-minute cache write", usage: dto.ClaudeUsage{CacheCreationInputTokens: 100, CacheCreation: &dto.ClaudeCacheCreationUsage{Ephemeral5mInputTokens: 100}}, wantQuota: 313},
|
|
{name: "cache read", usage: dto.ClaudeUsage{CacheReadInputTokens: 100}, wantQuota: 25},
|
|
{name: "no usage", usage: dto.ClaudeUsage{}, wantQuota: 0},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
info := &relaycommon.RelayInfo{
|
|
RelayFormat: types.RelayFormatClaude,
|
|
OriginModelName: "claude-opus-4-8",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 2.5, CompletionRatio: 5, CacheRatio: 0.1,
|
|
CacheCreationRatio: 1.25, CacheCreation5mRatio: 1.25, CacheCreation1hRatio: 2,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
}
|
|
usage := &dto.Usage{UsageSemantic: dto.BillingUsageSemanticAnthropic, BillingUsage: dto.NewClaudeMessagesBillingUsage(&tc.usage)}
|
|
summary := calculateTextQuotaSummary(ctx, info, effectiveBillingUsage(usage))
|
|
assert.Zero(t, summary.PromptTokens)
|
|
assert.Zero(t, summary.CompletionTokens)
|
|
assert.Equal(t, tc.wantQuota, summary.Quota)
|
|
assert.Equal(t, tc.wantQuota > 0, summary.hasBillableUsage())
|
|
t.Logf("input=%d output=%d quota=%d", summary.PromptTokens, summary.CompletionTokens, summary.Quota)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryDoesNotApplyRequestMultipliersToToolSurcharge(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "o1",
|
|
PriceData: hosttypes.PriceData{
|
|
ModelRatio: 1,
|
|
CompletionRatio: 1,
|
|
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
|
},
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolWebSearchPreview: {CallCount: 1},
|
|
},
|
|
},
|
|
}
|
|
relayInfo.PriceData.AddOtherRatio("n", 3)
|
|
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{})
|
|
|
|
expected := decimal.NewFromFloat(10.0 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expected.Equal(summary.ToolCallSurchargeQuota))
|
|
assert.Equal(t, common.QuotaFromDecimal(expected), summary.Quota)
|
|
}
|
|
|
|
func TestCalculateTextToolCallSurchargeGeminiGoogleSearch(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "gemini-3.7-flash",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolGoogleSearch: {ToolName: dto.BuildInToolGoogleSearch, CallCount: 1},
|
|
},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{ModelName: "gemini-3.7-flash", GroupRatio: 1}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
expected := decimal.NewFromFloat(14.0 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expected.Equal(surcharge), "got %s want %s", surcharge, expected)
|
|
require.Len(t, summary.ToolSurchargeItems, 1)
|
|
assert.Equal(t, dto.BuildInToolGoogleSearch, summary.ToolSurchargeItems[0].Name)
|
|
assert.Equal(t, 1, summary.ToolSurchargeItems[0].Count)
|
|
assert.Equal(t, 14.0, summary.ToolSurchargeItems[0].Price)
|
|
}
|
|
|
|
func TestCalculateTextToolCallSurchargeGeminiFunctionCall(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
|
|
operation_setting.SetToolPriceForTest("gemini_surcharge_fn", 5.0)
|
|
t.Cleanup(func() {
|
|
operation_setting.DeleteToolPriceForTest("gemini_surcharge_fn")
|
|
})
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "gemini-2.5-flash",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
"gemini_surcharge_fn": {CallCount: 2},
|
|
},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{ModelName: "gemini-2.5-flash", GroupRatio: 1}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
expected := decimal.NewFromFloat(5.0 * 2 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expected.Equal(surcharge), "got %s want %s", surcharge, expected)
|
|
require.Len(t, summary.ToolSurchargeItems, 1)
|
|
assert.Equal(t, "gemini_surcharge_fn", summary.ToolSurchargeItems[0].Name)
|
|
assert.Equal(t, 2, summary.ToolSurchargeItems[0].Count)
|
|
assert.Equal(t, 5.0, summary.ToolSurchargeItems[0].Price)
|
|
|
|
other := model.NewLogOther()
|
|
appendToolSurchargeLogInfo(other, summary.ToolSurchargeItems)
|
|
assert.Equal(t, summary.ToolSurchargeItems, other.Snapshot()["tool_surcharges"])
|
|
}
|
|
|
|
func TestCalculateTextToolCallSurchargeImageGenerationDefaultPrice(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
t.Cleanup(func() {
|
|
operation_setting.DeleteToolPriceForTest(dto.BuildInToolImageGeneration)
|
|
})
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "gpt-5.1",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolImageGeneration: {CallCount: 2},
|
|
},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{ModelName: "gpt-5.1", GroupRatio: 1.5}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
expected := decimal.NewFromFloat(150.0).
|
|
Mul(decimal.NewFromInt(2)).
|
|
Div(decimal.NewFromInt(1000)).
|
|
Mul(decimal.NewFromFloat(1.5)).
|
|
Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expected.Equal(surcharge), "got %s want %s", surcharge, expected)
|
|
require.Len(t, summary.ToolSurchargeItems, 1)
|
|
assert.Equal(t, dto.BuildInToolImageGeneration, summary.ToolSurchargeItems[0].Name)
|
|
assert.Equal(t, 2, summary.ToolSurchargeItems[0].Count)
|
|
assert.Equal(t, 150.0, summary.ToolSurchargeItems[0].Price)
|
|
}
|
|
|
|
func TestCalculateTextToolCallSurchargeImageGenerationExplicitZeroDisables(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
operation_setting.SetToolPriceForTest(dto.BuildInToolImageGeneration, 0)
|
|
t.Cleanup(func() {
|
|
operation_setting.DeleteToolPriceForTest(dto.BuildInToolImageGeneration)
|
|
})
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "gpt-5.1",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolImageGeneration: {CallCount: 3},
|
|
},
|
|
},
|
|
}
|
|
summary := &textQuotaSummary{ModelName: "gpt-5.1", GroupRatio: 1}
|
|
|
|
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
|
assert.True(t, surcharge.IsZero())
|
|
assert.Empty(t, summary.ToolSurchargeItems)
|
|
}
|
|
|
|
func TestCalculateTextQuotaSummaryImageGenerationUsesStructuredSurcharge(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
t.Cleanup(func() {
|
|
operation_setting.DeleteToolPriceForTest(dto.BuildInToolImageGeneration)
|
|
})
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
OriginModelName: "gpt-5.1",
|
|
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
|
dto.BuildInToolImageGeneration: {CallCount: 1},
|
|
},
|
|
},
|
|
}
|
|
relayInfo.PriceData.GroupRatioInfo.GroupRatio = 1
|
|
relayInfo.PriceData.ModelRatio = 1
|
|
relayInfo.PriceData.CompletionRatio = 1
|
|
|
|
usage := &dto.Usage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15}
|
|
summary := calculateTextQuotaSummary(ctx, relayInfo, usage)
|
|
|
|
require.Len(t, summary.ToolSurchargeItems, 1)
|
|
assert.Equal(t, dto.BuildInToolImageGeneration, summary.ToolSurchargeItems[0].Name)
|
|
assert.Equal(t, 1, summary.ToolSurchargeItems[0].Count)
|
|
assert.Equal(t, 150.0, summary.ToolSurchargeItems[0].Price)
|
|
|
|
expectedSurcharge := decimal.NewFromFloat(150.0 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
|
assert.True(t, expectedSurcharge.Equal(summary.ToolCallSurchargeQuota),
|
|
"got %s want %s", summary.ToolCallSurchargeQuota, expectedSurcharge)
|
|
assert.Greater(t, summary.Quota, 0)
|
|
}
|
|
|
|
func TestAppendToolSurchargeLogInfoWritesOnlyStructuredFields(t *testing.T) {
|
|
items := []ToolSurchargeItem{
|
|
{Name: dto.BuildInToolWebSearch, Count: 2, Price: 10},
|
|
{Name: dto.BuildInToolImageGeneration, Count: 1, Price: 150},
|
|
}
|
|
other := model.NewLogOther()
|
|
|
|
appendToolSurchargeLogInfo(other, items)
|
|
|
|
fields := other.Snapshot()
|
|
assert.Equal(t, items, fields["tool_surcharges"])
|
|
assert.NotContains(t, fields, "web_search")
|
|
assert.NotContains(t, fields, "web_search_call_count")
|
|
assert.NotContains(t, fields, "web_search_price")
|
|
assert.NotContains(t, fields, "file_search")
|
|
assert.NotContains(t, fields, "image_generation_call")
|
|
assert.NotContains(t, fields, "image_generation_call_price")
|
|
}
|
|
|
|
// A request whose upstream returns no usage is logged without a charge. The
|
|
// reason is stored as content_parts, which the web console renders in the
|
|
// viewer's language; content keeps the English text.
|
|
func TestPostTextConsumeQuotaLogsMissingUsageAsStructuredContent(t *testing.T) {
|
|
truncate(t)
|
|
const userID, channelID = 62, 62
|
|
seedUser(t, userID, 10_000)
|
|
seedChannel(t, channelID)
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
info := &relaycommon.RelayInfo{
|
|
UserId: userID,
|
|
OriginModelName: "missing-usage-model",
|
|
UsingGroup: "default",
|
|
StartTime: time.Now(),
|
|
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID},
|
|
PriceData: hosttypes.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1}},
|
|
}
|
|
|
|
PostTextConsumeQuota(ctx, info, &dto.Usage{}, nil)
|
|
|
|
log := getLastLog(t)
|
|
require.NotNil(t, log)
|
|
assert.Zero(t, log.Quota)
|
|
assert.Equal(t, "Upstream returned no usage, so nothing was charged (possibly an upstream timeout)", log.Content)
|
|
var other map[string]any
|
|
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
|
assert.Equal(t, []any{
|
|
map[string]any{"key": "Upstream returned no usage, so nothing was charged (possibly an upstream timeout)"},
|
|
}, other["content_parts"])
|
|
}
|